from __future__ import print_function
import argparse
import wandb
import os
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', required=True, help='cifar10 | lsun | mnist |imagenet | folder | lfw | fake')
parser.add_argument('--model', required=True, help='gramnet | gan')
parser.add_argument('--dataroot', required=True, help='path to dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=2)
parser.add_argument('--batchSize', type=int, default=100, help='input batch size')
parser.add_argument('--showSize', type=int, default=100, help='size of display batch')
parser.add_argument('--imageSize', type=int, default=64, help='the height / width of the input image to network')
parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector')
parser.add_argument('--nk', type=int, default=200, help='size of the projected k vector')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ncf', type=int, default=64)
parser.add_argument('--n_epochs', type=int, default=25, help='number of epochs to train for')
parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--clip_ratio', action='store_true', help='apply ratio clipping (suggested by reviewer 1)')
parser.add_argument('--eps_ratio', type=float, default=0.001, help='add eps to the diagonal before solving')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--gpu_id', type=int, default=0, help='default GPU ID to use')
parser.add_argument('--ngpu', type=int, default=1, help='number of GPUs to use')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--netF', default='', help="path to netF (to continue training)")
parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')
parser.add_argument('--monitor_heuristic', action='store_true', help='monitor heuristic σ')
parser.add_argument('--manualSeed', type=int, help='manual seed')
parser.add_argument('--nowandb', action='store_true', help='disables wandb')
parser.add_argument('--classes', default='bedroom', help='comma separated list of classes for the lsun data set')

opt = parser.parse_args()
print(opt)
if not opt.nowandb:
    wandb.init(project="gramtorch", config=opt)

try:
    os.makedirs(opt.outf)
except OSError:
    pass

if opt.manualSeed is None:
    opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)

cudnn.benchmark = True

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

if opt.dataset in ['imagenet', 'folder', 'lfw']:
    # folder dataset
    dataset = dset.ImageFolder(root=opt.dataroot,
                               transform=transforms.Compose([
                                   transforms.Resize(opt.imageSize),
                                   transforms.CenterCrop(opt.imageSize),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                               ]))
    nc=3
elif opt.dataset == 'lsun':
    classes = [ c + '_train' for c in opt.classes.split(',')]
    dataset = dset.LSUN(root=opt.dataroot, classes=classes,
                        transform=transforms.Compose([
                            transforms.Resize(opt.imageSize),
                            transforms.CenterCrop(opt.imageSize),
                            transforms.ToTensor(),
                            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                        ]))
    nc=3
elif opt.dataset == 'cifar10':
    dataset = dset.CIFAR10(root=opt.dataroot, download=True,
                           transform=transforms.Compose([
                               transforms.Resize(opt.imageSize),
                               transforms.ToTensor(),
                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                           ]))
    nc=3

elif opt.dataset == 'mnist':
        dataset = dset.MNIST(root=opt.dataroot, download=True,
                           transform=transforms.Compose([
                               transforms.Resize(opt.imageSize),
                               transforms.ToTensor(),
                               transforms.Normalize((0.5,), (0.5,)),
                           ]))
        nc=1

elif opt.dataset == 'fake':
    dataset = dset.FakeData(image_size=(3, opt.imageSize, opt.imageSize),
                            transform=transforms.ToTensor())
    nc=3

assert dataset
dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize,
                                         shuffle=True, num_workers=int(opt.workers))

device = torch.device(f"cuda:{opt.gpu_id}" if opt.cuda else "cpu")
ngpu = int(opt.ngpu)
nz = int(opt.nz)
nk = int(opt.nk)
ngf = int(opt.ngf)
ncf = int(opt.ncf)

# custom weights initialization called on netG and netD
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

def save_img(x_data, x_gen, epoch):
    vutils.save_image(
        x_data[0:opt.showSize],
        f'{opt.outf}/{opt.dataset}-{opt.model}-data.png',
        normalize=True
    )
    vutils.save_image(
        x_gen[0:opt.showSize],
        f'{opt.outf}/{opt.dataset}-{opt.model}-samples-epoch={epoch:03d}.png',
        normalize=True
    )
    if not opt.nowandb:
        wandb.log({"samples" : [wandb.Image(i) for i in x_gen[0:opt.showSize]]}, commit=False)

class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d(ngf * 2,     ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d(    ngf,      nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        if input.is_cuda and self.ngpu > 1:
            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
        else:
            output = self.main(input)
        return output

class Critic(nn.Module):
    def __init__(self, ngpu, nout):
        super(Critic, self).__init__()
        self.ngpu = ngpu
        self.nout = nout
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(nc, ncf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ncf) x 32 x 32
            nn.Conv2d(ncf, ncf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ncf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ncf*2) x 16 x 16
            nn.Conv2d(ncf * 2, ncf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ncf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ncf*4) x 8 x 8
            nn.Conv2d(ncf * 4, ncf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ncf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ncf*8) x 4 x 4
        )
        final = nn.Linear(ncf * 8 * 4 * 4, nout)
        if nout == 1:
            final = nn.Sequential(final, nn.Sigmoid())
        self.final = final

    def forward(self, input):
        if input.is_cuda and self.ngpu > 1:
            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
            output = output.view(-1, ncf * 8 * 4 * 4)
            output = nn.parallel.data_parallel(self.final, output, range(self.ngpu))
        else:
            output = self.main(input)
            output = output.view(-1, ncf * 8 * 4 * 4)
            output = self.final(output)
        if self.nout == 1:
            output = output.view(-1, 1).squeeze(1)
        return output

class GAN:
    def __init__(self, ngpu):
        netG = Generator(ngpu).to(device)
        netG.apply(weights_init)
        if opt.netG != '':
            netG.load_state_dict(torch.load(opt.netG))
        print(netG)

        netD = Critic(ngpu, 1).to(device)
        netD.apply(weights_init)
        if opt.netD != '':
            netD.load_state_dict(torch.load(opt.netD))
        print(netD)

        self.netG = netG
        self.netD = netD

    def train(self):
        netG = self.netG
        netD = self.netD

        criterion = nn.BCELoss()

        fixed_noise = torch.randn(opt.showSize, nz, 1, 1, device=device)
        real_label = 1
        fake_label = 0

        # setup optimizer
        optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
        optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))

        for epoch in range(opt.n_epochs):
            for i, data in enumerate(dataloader, 0):
                ############################
                # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
                ###########################
                # train with real
                netD.zero_grad()
                real_cpu = data[0].to(device)
                batch_size = real_cpu.size(0)
                label = torch.full((batch_size,), real_label, device=device)

                output = netD(real_cpu)
                errD_real = criterion(output, label)
                errD_real.backward()
                D_x = output.mean().item()

                # train with fake
                noise = torch.randn(batch_size, nz, 1, 1, device=device)
                fake = netG(noise)
                label.fill_(fake_label)
                output = netD(fake.detach())
                errD_fake = criterion(output, label)
                errD_fake.backward()
                D_G_z1 = output.mean().item()
                errD = errD_real + errD_fake
                optimizerD.step()

                ############################
                # (2) Update G network: maximize log(D(G(z)))
                ###########################
                netG.zero_grad()
                label.fill_(real_label)  # fake labels are real for generator cost
                output = netD(fake)
                errG = criterion(output, label)
                errG.backward()
                D_G_z2 = output.mean().item()
                optimizerG.step()

                print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'
                    % (epoch, opt.n_epochs, i, len(dataloader),
                        errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
                if i % 100 == 0:
                    netG.train(False)
                    save_img(real_cpu, netG(fixed_noise).detach(), epoch)
                    netG.train(True)
                if not opt.nowandb:
                    wandb.log({
                        "lossG" : errG, "lossF" : errD,
                        "D(x)" : D_G_z1, "D(G(z))" : D_G_z2,
                    })

            # do checkpointing
            torch.save(netG.state_dict(), f'{opt.outf}/{opt.model}-netG-epoch={epoch}.pth')
            torch.save(netD.state_dict(), f'{opt.outf}/{opt.model}-netD-epoch={epoch}.pth')

### MMD utilities

def euclidsq(x, y):
    return torch.pow(torch.cdist(x, y), 2)

def gaussian_gramian(esq, σ):
    return torch.exp(torch.div(-esq, 2 * σ**2))

def prepare(x_de, x_nu):
    return euclidsq(x_de, x_de), euclidsq(x_de, x_nu), euclidsq(x_nu, x_nu)

USE_SOLVE = True

def kmm_ratios(Kdede, Kdenu, λ):
    n_de, n_nu = Kdenu.shape
    if λ > 0:
        A = Kdede + λ * torch.eye(n_de).to(device)
    else:
        A = Kdede
    # Equivalent implement based on 1) solver and 2) matrix inversion
    if USE_SOLVE:
        B = torch.sum(Kdenu, 1, keepdim=True)
        return (n_de / n_nu) * torch.solve(B, A).solution
    else:
        B = Kdenu
        return torch.matmul(torch.matmul(torch.inverse(A), B), torch.ones(n_nu, 1).to(device))

def mmdsq_of(Kdede, Kdenu, Knunu):
    return torch.mean(Kdede) - 2 * torch.mean(Kdenu) + torch.mean(Knunu)

def estimate_ratio_compute_mmd(x_de, x_nu, σs):
    dsq_dede, dsq_denu, dsq_nunu = prepare(x_de, x_nu)
    if len(σs) == 0:
        # A heuristic is to use the median of pairwise distances as σ, suggested by Sugiyama's book
        sigma = torch.sqrt(
            torch.median(
                torch.cat([dsq_dede.squeeze(), dsq_denu.squeeze(), dsq_nunu.squeeze()], 1)
            )
        ).item()
        if not opt.nowandb:
            wandb.log({"heuristic_sigma" : sigma})
        elif opt.monitor_heuristic:
            print("heuristic sigma: ", sigma)
        # Use [sigma / 5, sigma / 3, sigma, sigma * 3, sigma * 5] if nothing provided
        if len(σs) == 0:
            σs.append(sigma)
            σs.append(sigma * 0.333)
            σs.append(sigma * 0.2)
            σs.append(sigma / 0.2)
            σs.append(sigma / 0.333)

    is_first = True
    ratio = None
    mmdsq = None
    for σ in σs:
        Kdede = gaussian_gramian(dsq_dede, σ)
        Kdenu = gaussian_gramian(dsq_denu, σ)
        Knunu = gaussian_gramian(dsq_nunu, σ)
        if is_first:
            ratio = kmm_ratios(Kdede, Kdenu, opt.eps_ratio)
            mmdsq = mmdsq_of(Kdede, Kdenu, Knunu)
            is_first = False
        else:
            ratio += kmm_ratios(Kdede, Kdenu, opt.eps_ratio)
            mmdsq += mmdsq_of(Kdede, Kdenu, Knunu)

    ratio = ratio / len(σs)
    ratio = torch.relu(ratio) if opt.clip_ratio else ratio
    mmd = torch.sqrt(torch.relu(mmdsq))

    return ratio, mmd

def extract_grad(m):
    gs = []
    for p in m.parameters():
        gs.append(p.grad.clone())
    return gs

def assign_grad(m, gs):
    for p, g in zip(m.parameters(), gs):
        p.grad = g

def sim_step(optimizer1, optimizer2, m1, m2, loss1, loss2):
    loss1.backward(retain_graph=True)
    gs1 = extract_grad(m1)
    m1.zero_grad()
    m2.zero_grad()
    loss2.backward()
    optimizer2.step()
    assign_grad(m1, gs1)
    optimizer1.step()

class GRAMnet:
    def __init__(self, ngpu):
        netG = Generator(ngpu).to(device)
        netG.apply(weights_init)
        if opt.netG != '':
            netG.load_state_dict(torch.load(opt.netG))
        print(netG)

        netF = Critic(ngpu, nk).to(device)
        netF.apply(weights_init)
        if opt.netF != '':
            netF.load_state_dict(torch.load(opt.netF))
        print(netF)

        self.netG = netG
        self.netF = netF

        if opt.dataset == "mnist":
            sigma_list = np.sqrt([0.01, 1, 100, 10000])
        elif opt.dataset == "cifar10":
            sigma_list = [1, 10, 100, 1000]
        else:
            sigma_list = [] # this will trigger automatic choice of sigma
        self.sigma_list = sigma_list

    def train(self):

        netG = self.netG
        netF = self.netF

        fixed_noise = torch.rand(opt.showSize, nz, 1, 1, device=device)

        # setup optimizer
        optimizerF = optim.Adam(netF.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
        optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))

        for epoch in range(opt.n_epochs):
            for i, data in enumerate(dataloader, 0):
                x_data = data[0].to(device)
                batch_size = x_data.size(0)

                netF.zero_grad()
                netG.zero_grad()
                # Generate samples
                noise = torch.rand(batch_size, nz, 1, 1, device=device)
                x_gen = netG(noise)
                # Project to low-dimensional space
                fx_data = netF(x_data)
                fx_gen = netF(x_gen)
                # Compute ratio and mmd
                ratio, mmd = estimate_ratio_compute_mmd(
                    fx_gen, fx_data, list(self.sigma_list)  # `list` is need to make a copy of sigmas
                )
                pearson_divergence = torch.mean(torch.pow(ratio - 1, 2))
                lossG = mmd
                lossF = -pearson_divergence
                # Add positivity regularizer if not clipping
                if not opt.clip_ratio:
                    lossF -= torch.sum(ratio)
                # Update G and F simultaneously
                sim_step(optimizerG, optimizerF, netG, netF, lossG, lossF)

                print('[%d/%d][%d/%d] Loss_F: %.4f Loss_G: %.4f'
                    % (epoch, opt.n_epochs, i, len(dataloader), lossF.item(), lossG.item()))
                if i % 100 == 0:
                    netG.train(False)
                    save_img(x_data, netG(fixed_noise).detach(), epoch)
                    netG.train(True)
                if not opt.nowandb:
                    wandb.log({
                        "lossG" : lossG, "lossF" : lossF, "mmd" : mmd,
                        "pearson_divergence" : pearson_divergence,
                    })

            # do checkpointing
            torch.save(netG.state_dict(), f'{opt.outf}/{opt.model}-netG-epoch={epoch}.pth')
            torch.save(netF.state_dict(), f'{opt.outf}/{opt.model}-netF-epoch={epoch}.pth')

if opt.model == "gan":
    model = GAN(ngpu)
elif opt.model == "gramnet":
    model = GRAMnet(ngpu)

model.train()


# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# (C) British Crown Copyright 2017-2021 Met Office.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
"""Reliability calibration plugins."""

import operator
import warnings

import iris
import numpy as np
import scipy

from improver import BasePlugin, PostProcessingPlugin
from improver.calibration.utilities import (
    check_forecast_consistency,
    create_unified_frt_coord,
    filter_non_matching_cubes,
)
from improver.metadata.probabilistic import (
    find_threshold_coordinate,
    probability_is_above_or_below,
)
from improver.metadata.utilities import generate_mandatory_attributes
from improver.utilities.cube_manipulation import MergeCubes, collapsed

class ConstructReliabilityCalibrationTables(BasePlugin):

    """A plugin for creating and populating reliability calibration tables."""

    def __init__(
        self,
        n_probability_bins=5,
        single_value_lower_limit=False,
        single_value_upper_limit=False,
    ):
        """
        Initialise class for creating reliability calibration tables. These
        tables include data columns entitled observation_count,
        sum_of_forecast_probabilities, and forecast_count, defined below.

        n_probability_bins (int):
            The total number of probability bins required in the reliability
            tables. If single value limits are turned on, these are included in
            this total.
        single_value_lower_limit (bool):
            Mandates that the lowest bin should be single valued,
            with a small precision tolerance, defined as 1.0E-6.
            The bin is thus 0 to 1.0E-6.
        single_value_upper_limit (bool):
            Mandates that the highest bin should be single valued,
            with a small precision tolerance, defined as 1.0E-6.
            The bin is thus (1 - 1.0E-6) to 1.
        """
        self.single_value_tolerance = 1.0e-6
        self.probability_bins = self._define_probability_bins(
            n_probability_bins, single_value_lower_limit, single_value_upper_limit
        )
        self.table_columns = np.array(
            ["observation_count", "sum_of_forecast_probabilities", "forecast_count"]
        )
        self.expected_table_shape = (len(self.table_columns), n_probability_bins)

    def __repr__(self):
        """Represent the configured plugin instance as a string."""
        bin_values = ", ".join(
            ["[{:1.2f} --> {:1.2f}]".format(*item) for item in self.probability_bins]
        )
        result = "<ConstructReliabilityCalibrationTables: " "probability_bins: {}>"
        return result.format(bin_values)

    def _define_probability_bins(
        self, n_probability_bins, single_value_lower_limit, single_value_upper_limit
    ):
        """
        Define equally sized probability bins for use in a reliability table.
        The range 0 to 1 is divided into ranges to give n_probability bins.
        If single_value_lower_limit and / or single_value_upper_limit are True,
        additional bins corresponding to values of 0 and / or 1 will be created,
        each with a width defined by self.single_value_tolerance.

        Args:
            n_probability_bins (int):
                The total number of probability bins desired in the
                reliability tables. This number includes the extrema bins
                (equals 0 and equals 1) if single value limits are turned on,
                in which case the minimum number of bins is 3.
            single_value_lower_limit (bool):
                Mandates that the lowest bin should be single valued,
                with a small precision tolerance, defined as 1.0E-6.
                The bin is thus 0 to 1.0E-6.
            single_value_upper_limit (bool):
                Mandates that the highest bin should be single valued,
                with a small precision tolerance, defined as 1.0E-6.
                The bin is thus (1 - 1.0E-6) to 1.
        Returns:
            numpy.ndarray:
                An array of 2-element arrays that contain the bounds of the
                probability bins. These bounds are non-overlapping, with
                adjacent bin boundaries spaced at the smallest representable
                interval.
        Raises:
            ValueError: If trying to use both single_value_lower_limit and
                        single_value_upper_limit with 2 or fewer probability bins.
        """
        if single_value_lower_limit and single_value_upper_limit:
            if n_probability_bins <= 2:
                msg = (
                    "Cannot use both single_value_lower_limit and "
                    "single_value_upper_limit with 2 or fewer "
                    "probability bins."
                )
                raise ValueError(msg)
            n_probability_bins = n_probability_bins - 2
        elif single_value_lower_limit or single_value_upper_limit:
            n_probability_bins = n_probability_bins - 1

        bin_lower = np.linspace(0, 1, n_probability_bins + 1, dtype=np.float32)
        bin_upper = np.nextafter(bin_lower, 0, dtype=np.float32)
        bin_upper[-1] = 1.0
        bins = np.stack([bin_lower[:-1], bin_upper[1:]], 1).astype(np.float32)

        if single_value_lower_limit:
            bins[0, 0] = np.nextafter(self.single_value_tolerance, 1, dtype=np.float32)
            lowest_bin = np.array([0, self.single_value_tolerance], dtype=np.float32)
            bins = np.vstack([lowest_bin, bins]).astype(np.float32)

        if single_value_upper_limit:
            bins[-1, 1] = np.nextafter(
                1.0 - self.single_value_tolerance, 0, dtype=np.float32
            )
            highest_bin = np.array(
                [1.0 - self.single_value_tolerance, 1], dtype=np.float32
            )
            bins = np.vstack([bins, highest_bin]).astype(np.float32)

        return bins

    def _create_probability_bins_coord(self):
        """
        Construct a dimension coordinate describing the probability bins
        of the reliability table.

        Returns:
            iris.coords.DimCoord:
                A dimension coordinate describing probability bins.
        """
        values = np.mean(self.probability_bins, axis=1, dtype=np.float32)
        probability_bins_coord = iris.coords.DimCoord(
            values, long_name="probability_bin", units=1, bounds=self.probability_bins
        )
        return probability_bins_coord

    def _create_reliability_table_coords(self):
        """
        Construct coordinates that describe the reliability table rows. These
        are observation_count, sum_of_forecast_probabilities, and
        forecast_count. The order used here is the order in which the table
        data is populated, so these must remain consistent with the
        _populate_reliability_bins function.

        Returns:
            (tuple): tuple containing:
                **index_coord** (iris.coords.DimCoord):
                    A numerical index dimension coordinate.
                **name_coord** (iris.coords.AuxCoord):
                    An auxiliary coordinate that assigns names to the index
                    coordinates, where these names correspond to the
                    reliability table rows.
        """
        index_coord = iris.coords.DimCoord(
            np.arange(len(self.table_columns), dtype=np.int32),
            long_name="table_row_index",
            units=1,
        )
        name_coord = iris.coords.AuxCoord(
            self.table_columns, long_name="table_row_name", units=1
        )
        return index_coord, name_coord

    @staticmethod
    def _define_metadata(forecast_slice):
        """
        Define metadata that is specifically required for reliability table
        cubes, whilst ensuring any mandatory attributes are also populated.

        Args:
            forecast_slice (iris.cube.Cube):
                The source cube from which to get pre-existing metadata of use.
        Returns:
            dict:
                A dictionary of attributes that are appropriate for the
                reliability table cube.
        """
        attributes = generate_mandatory_attributes([forecast_slice])
        attributes["title"] = "Reliability calibration data table"
        return attributes

    def _create_reliability_table_cube(self, forecast, threshold_coord):
        """
        Construct a reliability table cube and populate it with the provided
        data. The returned cube will include a cycle hour coordinate, which
        describes the model cycle hour at which the forecast data was produced.
        It will further include the forecast period, threshold coordinate,
        and spatial coordinates from the forecast cube.

        Args:
            forecast (iris.cube.Cube):
                A cube slice across the spatial dimensions of the forecast
                data. This slice provides the time and threshold values that
                relate to the reliability_table_data.
            threshold_coord (iris.coords.DimCoord):
                The threshold coordinate.
        Returns:
            iris.cube.Cube:
                A reliability table cube.
        """

        def _get_coords_and_dims(coord_names):
            """Obtain the requested coordinates and their dimension index from
            the forecast slice cube."""
            coords_and_dims = []
            leading_coords = [probability_bins_coord, reliability_index_coord]
            for coord_name in coord_names:
                crd = forecast_slice.coord(coord_name)
                crd_dim = forecast_slice.coord_dims(crd)
                crd_dim = crd_dim[0] + len(leading_coords) if crd_dim else ()
                coords_and_dims.append((crd, crd_dim))
            return coords_and_dims

        forecast_slice = next(forecast.slices_over(["time", threshold_coord]))
        expected_shape = self.expected_table_shape + forecast_slice.shape
        dummy_data = np.zeros((expected_shape))

        diagnostic = find_threshold_coordinate(forecast).name()
        attributes = self._define_metadata(forecast)

        # Define reliability table specific coordinates
        probability_bins_coord = self._create_probability_bins_coord()
        (
            reliability_index_coord,
            reliability_name_coord,
        ) = self._create_reliability_table_coords()
        frt_coord = create_unified_frt_coord(forecast.coord("forecast_reference_time"))

        # List of required non-spatial coordinates from the forecast
        non_spatial_coords = ["forecast_period", diagnostic]

        # Construct a list of coordinates in the desired order
        dim_coords = [forecast.coord(axis=dim).name() for dim in ["x", "y"]]
        dim_coords_and_dims = _get_coords_and_dims(dim_coords)
        aux_coords_and_dims = _get_coords_and_dims(non_spatial_coords)
        dim_coords_and_dims.append((reliability_index_coord, 0))
        aux_coords_and_dims.append((reliability_name_coord, 0))
        dim_coords_and_dims.append((probability_bins_coord, 1))

        reliability_cube = iris.cube.Cube(
            dummy_data,
            units=1,
            attributes=attributes,
            dim_coords_and_dims=dim_coords_and_dims,
            aux_coords_and_dims=aux_coords_and_dims,
        )
        reliability_cube.add_aux_coord(frt_coord)
        reliability_cube.rename("reliability_calibration_table")

        return reliability_cube

    def _populate_reliability_bins(self, forecast, truth):
        """
        For an x-y slice at a single validity time and threshold, populate
        a reliability table using the provided truth.

        Args:
            forecast (numpy.ndarray or numpy.ma.MaskedArray):
                An array containing data over an xy slice for a single validity
                time and threshold.
            truth (numpy.ndarray or numpy.ma.MaskedArray):
                An array containing a thresholded gridded truth at an
                equivalent validity time to the forecast array.
        Returns:
            numpy.ma.MaskedArray:
                An array containing reliability table data for a single time
                and threshold. The leading dimension corresponds to the rows
                of a calibration table, the second dimension to the number of
                probability bins, and the trailing dimensions are the spatial
                dimensions of the forecast and truth cubes (which are
                equivalent).
        """
        observation_counts = []
        forecast_probabilities = []
        forecast_counts = []

        for bin_min, bin_max in self.probability_bins:
            observation_mask = (
                ((forecast >= bin_min) & (forecast <= bin_max)) & (np.isclose(truth, 1))
            ).astype(int)
            forecast_mask = ((forecast >= bin_min) & (forecast <= bin_max)).astype(int)
            forecasts_probability_values = forecast * forecast_mask

            observation_counts.append(observation_mask)
            forecast_probabilities.append(forecasts_probability_values)
            forecast_counts.append(forecast_mask)

        reliability_table = np.ma.stack(
            [
                np.ma.stack(observation_counts),
                np.ma.stack(forecast_probabilities),
                np.ma.stack(forecast_counts),
            ]
        )

        return reliability_table.astype(np.float32)

    def _populate_masked_reliability_bins(self, forecast, truth):
        """
        Support populating the reliability table bins with a masked truth. If a
        masked truth is provided, a masked reliability table is returned.

        Args:
            forecast (numpy.ndarray):
                An array containing data over an xy slice for a single validity
                time and threshold.
            truth (numpy.ma.MaskedArray):
                An array containing a thresholded gridded truth at an
                equivalent validity time to the forecast array.
        Returns:
            numpy.ma.MaskedArray:
                An array containing reliability table data for a single time
                and threshold. The leading dimension corresponds to the rows
                of a calibration table, the second dimension to the number of
                probability bins, and the trailing dimensions are the spatial
                dimensions of the forecast and truth cubes (which are
                equivalent).
        """
        forecast = np.ma.masked_where(np.ma.getmask(truth), forecast)
        table = self._populate_reliability_bins(forecast, truth)
        # Zero data underneath mask to support bitwise addition of masks.
        table.data[table.mask] = 0
        return table

    def _add_reliability_tables(self, forecast, truth, threshold_reliability):
        """
        Add reliability tables. The presence of a masked truth is handled
        separately to ensure support for a mask that changes with validity time.

        Args:
            forecast (numpy.ndarray):
                An array containing data over an xy slice for a single validity
                time and threshold.
            truth (numpy.ndarray or numpy.ma.MaskedArray):
                An array containing a thresholded gridded truth at an
                equivalent validity time to the forecast array.
            threshold_reliability (numpy.ndarray or numpy.ma.MaskedArray):
                The current reliability table that will be added to.
        Returns:
            numpy.ndarray or numpy.ma.MaskedArray:
                An array containing reliability table data for a single time
                and threshold. The leading dimension corresponds to the rows
                of a calibration table, the second dimension to the number of
                probability bins, and the trailing dimensions are the spatial
                dimensions of the forecast and truth cubes (which are
                equivalent).
        """
        if np.ma.is_masked(truth.data):
            table = self._populate_masked_reliability_bins(forecast.data, truth.data)
            # Bitwise addition of masks. This ensures that only points that are
            # masked in both the existing and new reliability tables are kept
            # as being masked within the resulting reliability table.
            mask = threshold_reliability.mask & table.mask
            threshold_reliability = np.ma.array(
                threshold_reliability.data + table.data, mask=mask, dtype=np.float32,
            )
        else:
            np.add(
                threshold_reliability,
                self._populate_reliability_bins(forecast.data, truth.data),
                out=threshold_reliability,
                dtype=np.float32,
            )
        return threshold_reliability

    def process(self, historic_forecasts, truths):
        """
        Slice data over threshold and time coordinates to construct reliability
        tables. These are summed over time to give a single table for each
        threshold, constructed from all the provided historic forecasts and
        truths. If a masked truth is provided, a masked reliability table is
        returned. If the mask within the truth varies at different timesteps,
        any point that is unmasked for at least one timestep will have
        unmasked values within the reliability table. Therefore historic
        forecast points will only be used if they have a corresponding valid
        truth point for each timestep.

        .. See the documentation for an example of the resulting reliability
           table cube.
        .. include:: extended_documentation/calibration/
           reliability_calibration/reliability_calibration_examples.rst

        Note that the forecast and truth data used is probabilistic, i.e. has
        already been thresholded relative to the thresholds of interest, using
        the equality operator required. As such this plugin is agnostic as to
        whether the data is thresholded below or above a given diagnostic
        threshold.

        Args:
            historic_forecasts (iris.cube.Cube):
                A cube containing the historical forecasts used in calibration.
                These are expected to all have a consistent cycle hour, that is
                the hour in the forecast reference time.
            truths (iris.cube.Cube):
                A cube containing the thresholded gridded truths used in
                calibration.
        Returns:
            iris.cube.CubeList:
                A cubelist of reliability table cubes, one for each threshold
                in the historic forecast cubes.
        Raises:
            ValueError: If the forecast and truth cubes have differing
                        threshold coordinates.
        """
        historic_forecasts, truths = filter_non_matching_cubes(
            historic_forecasts, truths
        )

        threshold_coord = find_threshold_coordinate(historic_forecasts)
        truth_threshold_coord = find_threshold_coordinate(truths)
        if not threshold_coord == truth_threshold_coord:
            msg = "Threshold coordinates differ between forecasts and truths."
            raise ValueError(msg)

        time_coord = historic_forecasts.coord("time")

        check_forecast_consistency(historic_forecasts)
        reliability_cube = self._create_reliability_table_cube(
            historic_forecasts, threshold_coord
        )

        populate_bins_func = self._populate_reliability_bins
        if np.ma.is_masked(truths.data):
            populate_bins_func = self._populate_masked_reliability_bins

        reliability_tables = iris.cube.CubeList()
        threshold_slices = zip(
            historic_forecasts.slices_over(threshold_coord),
            truths.slices_over(threshold_coord),
        )
        for forecast_slice, truth_slice in threshold_slices:

            time_slices = zip(
                forecast_slice.slices_over(time_coord),
                truth_slice.slices_over(time_coord),
            )
            forecast, truth = next(time_slices)
            threshold_reliability = populate_bins_func(forecast.data, truth.data)

            for forecast, truth in time_slices:
                threshold_reliability = self._add_reliability_tables(
                    forecast, truth, threshold_reliability
                )

            reliability_entry = reliability_cube.copy(data=threshold_reliability)
            reliability_entry.replace_coord(forecast_slice.coord(threshold_coord))
            reliability_tables.append(reliability_entry)

        return MergeCubes()(reliability_tables, copy=False)

class AggregateReliabilityCalibrationTables(BasePlugin):

    """This plugin enables the aggregation of multiple reliability calibration
    tables, and/or the aggregation over coordinates in the tables."""

    def __repr__(self):
        """Represent the configured plugin instance as a string."""
        return "<AggregateReliabilityCalibrationTables>"

    @staticmethod
    def _check_frt_coord(cubes):
        """
        Check that the reliability calibration tables do not have overlapping
        forecast reference time bounds. If these coordinates overlap in time it
        indicates that some of the same forecast data has contributed to more
        than one table, thus aggregating them would double count these
        contributions.

        Args:
            cubes (iris.cube.CubeList):
                The list of reliability calibration tables for which the
                forecast reference time coordinates should be checked.
        Raises:
            ValueError: If the bounds overlap.
        """
        lower_bounds = []
        upper_bounds = []
        for cube in cubes:
            lower_bounds.append(cube.coord("forecast_reference_time").bounds[0][0])
            upper_bounds.append(cube.coord("forecast_reference_time").bounds[0][1])
        if not all(x < y for x, y in zip(upper_bounds, lower_bounds[1:])):
            raise ValueError(
                "Reliability calibration tables have overlapping "
                "forecast reference time bounds, indicating that "
                "the same forecast data has contributed to the "
                "construction of both tables. Cannot aggregate."
            )

    def process(self, cubes, coordinates=None):
        """
        Aggregate the input reliability calibration table cubes and return the
        result.

        Args:
            cubes (list or iris.cube.CubeList):
                The cube or cubes containing the reliability calibration tables
                to aggregate.
            coordinates (list or None):
                A list of coordinates over which to aggregate the reliability
                calibration table using summation. If the argument is None and
                a single cube is provided, this cube will be returned
                unchanged.
        """
        coordinates = [] if coordinates is None else coordinates

        try:
            (cube,) = cubes
        except ValueError:
            cubes = iris.cube.CubeList(cubes)
            self._check_frt_coord(cubes)
            cube = cubes.merge_cube()
            coordinates.append("forecast_reference_time")
        else:
            if not coordinates:
                return cube

        result = collapsed(cube, coordinates, iris.analysis.SUM)
        frt = create_unified_frt_coord(cube.coord("forecast_reference_time"))
        result.replace_coord(frt)
        return result

class ManipulateReliabilityTable(BasePlugin):
    """
    A plugin to manipulate the reliability tables before they are used to
    calibrate a forecast. x and y coordinates on the reliability table must be
    collapsed.
    The result is a reliability diagram with monotonic observation frequency.

    Steps taken are:

    1. If any bin contains less than the minimum forecast count then try
    combining this bin with whichever neighbour has the lowest sample count.
    This process is repeated for all bins that are below the minimum forecast
    count criterion.

    2. If non-monotonicity of the observation frequency is detected, try
    combining a pair of bins that appear non-monotonic. Only a single pair of
    bins are combined.

    3. If non-monotonicity of the observation frequency remains after trying
    to combine a single pair of bins, replace non-monotonic bins by assuming a
    constant observation frequency.
    """

    def __init__(self, minimum_forecast_count=200):
        """
        Initialise class for manipulating a reliability table.

        Args:
            minimum_forecast_count (int):
                The minimum number of forecast counts in a forecast probability
                bin for it to be used in calibration.
                The default value of 200 is that used in Flowerdew 2014.
        Raises:
            ValueError: If minimum_forecast_count is less than 1.
        References:
            Flowerdew J. 2014. Calibrating ensemble reliability whilst
            preserving spatial structure. Tellus, Ser. A Dyn. Meteorol.
            Oceanogr. 66.
        """
        if minimum_forecast_count < 1:
            raise ValueError(
                "The minimum_forecast_count must be at least 1 as empty "
                "bins in the reliability table are not handled."
            )

        self.minimum_forecast_count = minimum_forecast_count

    @staticmethod
    def _extract_reliability_table_components(reliability_table):
        """Extract reliability table components from cube

        Args:
            reliability_table (iris.cube.Cube):
                A reliability table to be manipulated.

        Returns:
            Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, iris.coords.DimCoord]:
                Tuple containing the updated observation count,
                forecast probability sum, forecast count and probability bin
                coordinate.
        """
        observation_count = reliability_table.extract(
            iris.Constraint(table_row_name="observation_count")
        ).data
        forecast_probability_sum = reliability_table.extract(
            iris.Constraint(table_row_name="sum_of_forecast_probabilities")
        ).data
        forecast_count = reliability_table.extract(
            iris.Constraint(table_row_name="forecast_count")
        ).data
        probability_bin_coord = reliability_table.coord("probability_bin")
        return (
            observation_count,
            forecast_probability_sum,
            forecast_count,
            probability_bin_coord,
        )

    @staticmethod
    def _sum_pairs(array, upper):
        """
        Returns a new array where a pair of values in the original array have
        been replaced by their sum. Combines the value in the upper index with
        the value in the upper-1 index.

        Args:
            array (numpy.ndarray):
                Array to be modified.
            upper (int):
                Upper index of pair.

        Returns:
            numpy.ndarray:
                Array where a pair of values has been replaced by their sum.
        """
        result = array.copy()
        result[upper - 1] = np.sum(array[upper - 1 : upper + 1])
        return np.delete(result, upper)

    @staticmethod
    def _create_new_bin_coord(probability_bin_coord, upper):
        """
        Create a new probability_bin coordinate by combining two adjacent
        points on the probability_bin coordinate. This matches the combination
        of the data for the two bins.

        Args:
            probability_bin_coord (iris.coords.DimCoord):
                Original probability bin coordinate.
            upper (int):
                Upper index of pair.

        Returns:
            iris.coords.DimCoord:
                Probability bin coordinate with updated points and bounds where
                a pair of bins have been combined to create a single bin.
        """
        old_bounds = probability_bin_coord.bounds
        new_bounds = np.concatenate(
            (
                old_bounds[0 : upper - 1],
                np.array([[old_bounds[upper - 1, 0], old_bounds[upper, 1]]]),
                old_bounds[upper + 1 :],
            )
        )
        new_points = np.mean(new_bounds, axis=1, dtype=np.float32)
        new_bin_coord = iris.coords.DimCoord(
            new_points, long_name="probability_bin", units=1, bounds=new_bounds
        )
        return new_bin_coord

    def _combine_undersampled_bins(
        self,
        observation_count,
        forecast_probability_sum,
        forecast_count,
        probability_bin_coord,
    ):
        """
        Combine bins that are under-sampled i.e. that have a lower forecast
        count than the minimum_forecast_count, so that information from these
        poorly-sampled bins can contribute to the calibration. If multiple
        bins are below the minimum forecast count, the bin closest to
        meeting the minimum_forecast_count criterion is combined with whichever
        neighbour has the lowest sample count. A new bin is then created by
        summing the neighbouring pair of bins. This process is repeated for all
        bins that are below the minimum forecast count criterion.

        Args:
            observation_count (numpy.ndarray):
                Observation count extracted from reliability table.
            forecast_probability_sum (numpy.ndarray):
                Forecast probability sum extracted from reliability table.
            forecast_count (numpy.ndarray):
                Forecast count extracted from reliability table.
            probability_bin_coord (iris.coords.DimCoord):
                Original probability bin coordinate.
        Returns:
            Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, iris.coords.DimCoord]
                Tuple containing the updated observation count,
                forecast probability sum, forecast count and probability bin
                coordinate.

        """
        while (
            any(x < self.minimum_forecast_count for x in forecast_count)
            and len(forecast_count) > 1
        ):
            forecast_count_copy = forecast_count.copy()

            # Find index of the bin with the highest forecast count that is
            # below the minimum_forecast_count by setting forecast counts
            # greater than the minimum_forecast_count to NaN.
            forecast_count_copy[forecast_count >= self.minimum_forecast_count] = np.nan
            # Note for multiple occurrences of the maximum,
            # the index of the first occurrence is returned.
            index = np.int32(np.nanargmax(forecast_count_copy))

            # Determine the upper index of the pair of bins to be combined.
            if index == 0:
                # Must use higher bin
                upper = index + 1
            elif index + 1 == len(forecast_count):
                # Index already defines the upper bin
                upper = index
            else:
                # Define upper index to include bin with lowest sample count.
                if forecast_count[index + 1] > forecast_count[index - 1]:
                    upper = index
                else:
                    upper = index + 1

            forecast_count = self._sum_pairs(forecast_count, upper)
            observation_count = self._sum_pairs(observation_count, upper)
            forecast_probability_sum = self._sum_pairs(forecast_probability_sum, upper)
            probability_bin_coord = self._create_new_bin_coord(
                probability_bin_coord, upper
            )

        return (
            observation_count,
            forecast_probability_sum,
            forecast_count,
            probability_bin_coord,
        )

    def _combine_bin_pair(
        self,
        observation_count,
        forecast_probability_sum,
        forecast_count,
        probability_bin_coord,
    ):
        """
        Combine a pair of bins when non-monotonicity of the observation
        frequency is detected. Iterate top-down from the highest forecast
        probability bin to the lowest probability bin when combining the bins.
        Only allow a single pair of bins to be combined.

        Args:
            observation_count (numpy.ndarray):
                Observation count extracted from reliability table.
            forecast_probability_sum (numpy.ndarray):
                Forecast probability sum extracted from reliability table.
            forecast_count (numpy.ndarray):
                Forecast count extracted from reliability table.
            probability_bin_coord (iris.coords.DimCoord):
                Original probability bin coordinate.

        Returns:
            Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, iris.coords.DimCoord]
                Tuple containing the updated observation count,
                forecast probability sum, forecast count and probability bin
                coordinate.
        """
        observation_frequency = np.array(observation_count / forecast_count)
        for upper in np.arange(len(observation_frequency) - 1, 0, -1):
            (diff,) = np.diff(
                [observation_frequency[upper - 1], observation_frequency[upper]]
            )
            if diff < 0:
                forecast_count = self._sum_pairs(forecast_count, upper)
                observation_count = self._sum_pairs(observation_count, upper)
                forecast_probability_sum = self._sum_pairs(
                    forecast_probability_sum, upper
                )
                probability_bin_coord = self._create_new_bin_coord(
                    probability_bin_coord, upper
                )
                break
        return (
            observation_count,
            forecast_probability_sum,
            forecast_count,
            probability_bin_coord,
        )

    @staticmethod
    def _assume_constant_observation_frequency(observation_count, forecast_count):
        """
        Decide which end bin (highest probability bin or lowest probability
        bin) has the highest sample count. Iterate through the observation
        frequency from the end bin with the highest sample count to the end bin
        with the lowest sample count. Whilst iterating, compare each pair of
        bins and, if a pair is non-monotonic, replace the value of the bin
        closer to the lowest sample count end bin with the value of the
        bin that is closer to the higher sample count end bin. Then calculate
        the new observation count required to give a monotonic observation
        frequency.

        Args:
            observation_count (numpy.ndarray):
                Observation count extracted from reliability table.
            forecast_count (numpy.ndarray):
                Forecast count extracted from reliability table.

        Returns:
            numpy.ndarray:
                Observation count computed from a monotonic observation frequency.

        """
        observation_frequency = np.array(observation_count / forecast_count)

        iterator = observation_frequency
        operation = operator.lt
        # Top down if forecast count is lower for lowest probability bin,
        # than for highest probability bin.
        if forecast_count[0] < forecast_count[-1]:
            # Reverse array to iterate from top to bottom.
            iterator = observation_frequency[::-1]
            operation = operator.gt

        for index, lower_bin in enumerate(iterator[:-1]):
            (diff,) = np.diff([lower_bin, iterator[index + 1]])
            if operation(diff, 0):
                iterator[index + 1] = lower_bin

        observation_frequency = iterator
        if forecast_count[0] < forecast_count[-1]:
            # Re-reverse array from bottom to top to ensure original ordering.
            observation_frequency = iterator[::-1]

        observation_count = observation_frequency * forecast_count
        return observation_count

    @staticmethod
    def _update_reliability_table(
        reliability_table,
        observation_count,
        forecast_probability_sum,
        forecast_count,
        probability_bin_coord,
    ):
        """
        Update the reliability table data and the probability bin coordinate.

        Args:
            reliability_table (iris.cube.Cube):
                A reliability table to be manipulated.
            observation_count (numpy.ndarray):
                Observation count extracted from reliability table.
            forecast_probability_sum (numpy.ndarray):
                Forecast probability sum extracted from reliability table.
            forecast_count (numpy.ndarray):
                Forecast count extracted from reliability table.
            probability_bin_coord (iris.coords.DimCoord):
                Original probability bin coordinate.

        Returns:
            iris.cube.Cube:
                Updated reliability table.
        """
        final_data = np.stack(
            [observation_count, forecast_probability_sum, forecast_count]
        )
        nrows, ncols = final_data.shape
        reliability_table = reliability_table[0:nrows, 0:ncols].copy(data=final_data)
        reliability_table.replace_coord(probability_bin_coord)
        return reliability_table

    def process(self, reliability_table):
        """
        Apply the steps needed to produce a reliability diagram with a
        monotonic observation frequency.

        Args:
            reliability_table (iris.cube.Cube):
                A reliability table to be manipulated. The only coordinates
                expected on this cube are a threshold coordinate,
                a table_row_index coordinate and corresponding table_row_name
                coordinate and a probability_bin coordinate.

        Returns:
            iris.cube.CubeList:
                Containing a reliability table cube for each threshold in the
                input reliablity table. For tables where monotonicity has been
                enforced the probability_bin coordinate will have one less
                bin than the tables that were already monotonic. If
                under-sampled bins have been combined, then the probability_bin
                coordinate will have been reduced until all bins have more than
                the minimum_forecast_count if possible; a single under-sampled
                bin will be returned if combining all bins is still insufficient
                to reach the minimum_forecast_count.
        """
        threshold_coord = find_threshold_coordinate(reliability_table)
        reliability_table_cubelist = iris.cube.CubeList()
        for rel_table_slice in reliability_table.slices_over(threshold_coord):
            (
                observation_count,
                forecast_probability_sum,
                forecast_count,
                probability_bin_coord,
            ) = self._extract_reliability_table_components(rel_table_slice)

            if np.any(forecast_count < self.minimum_forecast_count):
                (
                    observation_count,
                    forecast_probability_sum,
                    forecast_count,
                    probability_bin_coord,
                ) = self._combine_undersampled_bins(
                    observation_count,
                    forecast_probability_sum,
                    forecast_count,
                    probability_bin_coord,
                )
                rel_table_slice = self._update_reliability_table(
                    rel_table_slice,
                    observation_count,
                    forecast_probability_sum,
                    forecast_count,
                    probability_bin_coord,
                )

            # If the observation frequency is non-monotonic adjust the
            # reliability table
            observation_frequency = np.array(observation_count / forecast_count)
            if not np.all(np.diff(observation_frequency) >= 0):
                (
                    observation_count,
                    forecast_probability_sum,
                    forecast_count,
                    probability_bin_coord,
                ) = self._combine_bin_pair(
                    observation_count,
                    forecast_probability_sum,
                    forecast_count,
                    probability_bin_coord,
                )
                observation_count = self._assume_constant_observation_frequency(
                    observation_count, forecast_count
                )
                rel_table_slice = self._update_reliability_table(
                    rel_table_slice,
                    observation_count,
                    forecast_probability_sum,
                    forecast_count,
                    probability_bin_coord,
                )
            reliability_table_cubelist.append(rel_table_slice)
        return reliability_table_cubelist

class ApplyReliabilityCalibration(PostProcessingPlugin):

    """
    A plugin for the application of reliability calibration to probability
    forecasts. This calibration is designed to improve the reliability of
    probability forecasts without significantly degrading their resolution.

    The method implemented here is described in Flowerdew J. 2014. Calibration
    is always applied as long as there are at least two bins within the input
    reliability table.

    References:
    Flowerdew J. 2014. Calibrating ensemble reliability whilst
    preserving spatial structure. Tellus, Ser. A Dyn. Meteorol.
    Oceanogr. 66.
    """

    def __init__(self):
        """
        Initialise class for applying reliability calibration.

        """
        self.threshold_coord = None

    @staticmethod
    def _extract_matching_reliability_table(forecast, reliability_table):
        """
        Extract the reliability table with a threshold coordinate
        matching the forecast cube.
        If no matching reliability table is found raise an exception.

        Args:
            forecast (iris.cube.Cube):
                The forecast to be calibrated.
            reliability_table (iris.cube.CubeList):
                The reliability table to use for applying calibration.
        Returns:
            iris.cube.Cube:
                A reliability table who's threshold coordinate matches
                the forecast cube.
        Raises:
            ValueError: If no matching reliability table is found.
        """
        threshold_coord = find_threshold_coordinate(forecast)
        coord_values = {threshold_coord.name(): threshold_coord.points}
        constr = iris.Constraint(coord_values=coord_values)
        if isinstance(reliability_table, iris.cube.Cube):
            extracted = reliability_table.extract(constr)
        else:
            extracted = reliability_table.extract(constr, strict=True)
        if not extracted:
            raise ValueError(
                "No reliability table found to match threshold "
                f"{find_threshold_coordinate(forecast).points[0]}."
            )
        return extracted

    def _ensure_monotonicity_across_thresholds(self, cube):
        """
        Ensures that probabilities change monotonically relative to thresholds
        in the expected order, e.g. exceedance probabilities always remain the
        same or decrease as the threshold values increase, below threshold
        probabilities always remain the same or increase as the threshold
        values increase.

        Args:
            cube (iris.cube.Cube):
                The probability cube for which monotonicity is to be checked
                and enforced. This cube is modified in place.
        Raises:
            ValueError: Threshold coordinate lacks the
                        spp__relative_to_threshold attribute.
        Warns:
            UserWarning: If the probabilities must be sorted to reinstate
                         expected monotonicity following calibration.
        """
        (threshold_dim,) = cube.coord_dims(self.threshold_coord)
        thresholding = probability_is_above_or_below(cube)
        if thresholding is None:
            msg = (
                "Cube threshold coordinate does not define whether "
                "thresholding is above or below the defined thresholds."
            )
            raise ValueError(msg)

        if (
            thresholding == "above"
            and not (np.diff(cube.data, axis=threshold_dim) <= 0).all()
        ):
            msg = (
                "Exceedance probabilities are not decreasing monotonically "
                "as the threshold values increase. Forced back into order."
            )
            warnings.warn(msg)
            cube.data = np.sort(cube.data, axis=threshold_dim)[::-1]

        if (
            thresholding == "below"
            and not (np.diff(cube.data, axis=threshold_dim) >= 0).all()
        ):
            msg = (
                "Below threshold probabilities are not increasing "
                "monotonically as the threshold values increase. Forced "
                "back into order."
            )
            warnings.warn(msg)
            cube.data = np.sort(cube.data, axis=threshold_dim)

    def _calculate_reliability_probabilities(self, reliability_table):
        """
        Calculates forecast probabilities and observation frequencies from the
        reliability table. If fewer than two bins are provided, Nones are
        returned as no calibration can be applied. Fewer than two bins can occur
        due to repeated combination of undersampled probability bins,
        please see :class:`.ManipulateReliabilityTable`.

        Args:
            reliability_table (iris.cube.Cube):
                A reliability table for a single threshold from which to
                calculate the forecast probabilities and observation
                frequencies.
        Returns:
            Optional[Tuple[numpy.ndarray, numpy.ndarray]]:
                Tuple containing forecast probabilities calculated by dividing
                the sum of forecast probabilities by the forecast count and
                observation frequency calculated by dividing the observation
                count by the forecast count.
        """
        observation_count = reliability_table.extract(
            iris.Constraint(table_row_name="observation_count")
        ).data
        forecast_count = reliability_table.extract(
            iris.Constraint(table_row_name="forecast_count")
        ).data
        forecast_probability_sum = reliability_table.extract(
            iris.Constraint(table_row_name="sum_of_forecast_probabilities")
        ).data

        # If there are fewer than two bins, no calibration can be applied.
        if len(np.atleast_1d(forecast_count)) < 2:
            return None, None

        forecast_probability = np.array(forecast_probability_sum / forecast_count)
        observation_frequency = np.array(observation_count / forecast_count)

        return forecast_probability, observation_frequency

    @staticmethod
    def _interpolate(
        forecast_threshold, reliability_probabilities, observation_frequencies
    ):
        """
        Perform interpolation of the forecast probabilities using the
        reliability table data to produce the calibrated forecast. Where
        necessary linear extrapolation will be applied. Any mask in place on
        the forecast_threshold data is removed and reapplied after calibration.

        Args:
            forecast_threshold (numpy.ndarray):
                The forecast probabilities to be calibrated.
            reliability_probabilities (numpy.ndarray):
                Probabilities taken from the reliability tables.
            observation_frequencies (numpy.ndarray):
                Observation frequencies that relate to the reliability
                probabilities, taken from the reliability tables.

        Returns:
            numpy.ndarray:
                The calibrated forecast probabilities. The final results are
                clipped to ensure any extrapolation has not yielded
                probabilities outside the range 0 to 1.
        """
        shape = forecast_threshold.shape
        mask = forecast_threshold.mask if np.ma.is_masked(forecast_threshold) else None

        forecast_probabilities = np.ma.getdata(forecast_threshold).flatten()

        interpolation_function = scipy.interpolate.interp1d(
            reliability_probabilities, observation_frequencies, fill_value="extrapolate"
        )
        interpolated = interpolation_function(forecast_probabilities.data)

        interpolated = interpolated.reshape(shape).astype(np.float32)

        if mask is not None:
            interpolated = np.ma.masked_array(interpolated, mask=mask)

        return np.clip(interpolated, 0, 1)

    def process(self, forecast, reliability_table):
        """
        Apply reliability calibration to a forecast. The reliability table
        and the forecast cube must share an identical threshold coordinate.

        Args:
            forecast (iris.cube.Cube):
                The forecast to be calibrated.
            reliability_table (iris.cube.Cube or iris.cube.CubeList):
                The reliability table to use for applying calibration.
                x and y dimensions must be collapsed.
        Returns:
            iris.cube.Cube:
                The forecast cube following calibration.
        """
        self.threshold_coord = find_threshold_coordinate(forecast)

        forecast_thresholds = forecast.slices_over(self.threshold_coord)

        uncalibrated_thresholds = []
        calibrated_cubes = iris.cube.CubeList()
        for forecast_threshold in forecast_thresholds:
            reliability_threshold = self._extract_matching_reliability_table(
                forecast_threshold, reliability_table
            )
            (
                reliability_probabilities,
                observation_frequencies,
            ) = self._calculate_reliability_probabilities(reliability_threshold)

            if reliability_probabilities is None:
                calibrated_cubes.append(forecast_threshold)
                uncalibrated_thresholds.append(
                    forecast_threshold.coord(self.threshold_coord).points[0]
                )
                continue

            interpolated = self._interpolate(
                forecast_threshold.data,
                reliability_probabilities,
                observation_frequencies,
            )

            calibrated_cubes.append(forecast_threshold.copy(data=interpolated))

        calibrated_forecast = calibrated_cubes.merge_cube()
        self._ensure_monotonicity_across_thresholds(calibrated_forecast)

        if uncalibrated_thresholds:
            msg = (
                "The following thresholds were not calibrated due to "
                "insufficient forecast counts in reliability table bins: "
                "{}".format(uncalibrated_thresholds)
            )
            warnings.warn(msg)

        return calibrated_forecast


"""
Operation representation classes for the `statevec_slow` evolution type.
"""
#***************************************************************************************************
# 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 itertools as _itertools
import copy as _copy

import numpy as _np
from scipy.sparse.linalg import LinearOperator

from .statereps import StateRepDensePure as _StateRepDensePure
from .. import basereps as _basereps
from pygsti.baseobjs.statespace import StateSpace as _StateSpace
from ...tools import basistools as _bt
from ...tools import internalgates as _itgs
from ...tools import optools as _ot

class OpRep(_basereps.OpRep):
    def __init__(self, state_space):
        self.state_space = state_space

    @property
    def dim(self):
        return self.state_space.udim

    def acton(self, state):
        raise NotImplementedError()

    def adjoint_acton(self, state):
        raise NotImplementedError()

    def aslinearoperator(self):
        def mv(v):
            if v.ndim == 2 and v.shape[1] == 1: v = v[:, 0]
            in_state = _StateRepDensePure(_np.ascontiguousarray(v, complex), self.state_space, basis=None)
            return self.acton(in_state).to_dense('Hilbert')

        def rmv(v):
            if v.ndim == 2 and v.shape[1] == 1: v = v[:, 0]
            in_state = _StateRepDensePure(_np.ascontiguousarray(v, complex), self.state_space, basis=None)
            return self.adjoint_acton(in_state).to_dense('Hilbert')
        return LinearOperator((self.dim, self.dim), matvec=mv, rmatvec=rmv)  # transpose, adjoint, dot, matmat?

    def copy(self):
        return _copy.deepcopy(self)

class OpRepDenseUnitary(OpRep):
    def __init__(self, mx, basis, state_space):
        state_space = _StateSpace.cast(state_space)
        if mx is None:
            mx = _np.identity(state_space.udim, complex)
        assert(mx.ndim == 2 and mx.shape[0] == state_space.udim)
        self.basis = basis
        self.base = _np.require(mx, requirements=['OWNDATA', 'C_CONTIGUOUS'])
        super(OpRepDenseUnitary, self).__init__(state_space)

    def base_has_changed(self):
        pass

    def to_dense(self, on_space):
        if on_space in ('minimal', 'Hilbert'):
            return self.base
        elif on_space == 'HilbertSchmidt':
            return _ot.unitary_to_superop(self.base, self.basis)
        else:
            raise ValueError("Invalid `on_space` argument: %s" % str(on_space))

    def acton(self, state):
        return _StateRepDensePure(_np.dot(self.base, state.data), state.state_space, state.basis)

    def adjoint_acton(self, state):
        return _StateRepDensePure(_np.dot(_np.conjugate(self.base.T), state.data),
                                  state.state_space, state.basis)

    def __str__(self):
        return "OpRepDenseUnitary:\n" + str(self.base)

class OpRepStandard(OpRepDenseUnitary):
    def __init__(self, name, basis, state_space):
        std_unitaries = _itgs.standard_gatename_unitaries()
        self.name = name
        if self.name not in std_unitaries:
            raise ValueError("Name '%s' not in standard unitaries" % self.name)

        U = std_unitaries[self.name]
        state_space = _StateSpace.cast(state_space)
        assert(U.shape[0] == state_space.udim)

        super(OpRepStandard, self).__init__(U, basis, state_space)

#class OpRepStochastic(OpRepDense):
# - maybe we could add this, but it wouldn't be a "dense" op here,
#   perhaps we need to change API?

class OpRepComposed(OpRep):
    # exactly the same as densitymx case
    def __init__(self, factor_op_reps, state_space):
        #assert(len(factor_op_reps) > 0), "Composed gates must contain at least one factor gate!"
        self.factors_reps = factor_op_reps
        super(OpRepComposed, self).__init__(state_space)

    def acton(self, state):
        """ Act this gate map on an input state """
        for gate in self.factor_reps:
            state = gate.acton(state)
        return state

    def adjoint_acton(self, state):
        """ Act the adjoint of this operation matrix on an input state """
        for gate in reversed(self.factor_reps):
            state = gate.adjoint_acton(state)
        return state

    def reinit_factor_op_reps(self, new_factor_op_reps):
        self.factors_reps = new_factor_op_reps

class OpRepSum(OpRep):
    # exactly the same as densitymx case
    def __init__(self, factor_reps, state_space):
        #assert(len(factor_reps) > 0), "Composed gates must contain at least one factor gate!"
        self.factor_reps = factor_reps
        super(OpRepSum, self).__init__(state_space)

    def acton(self, state):
        """ Act this gate map on an input state """
        output_state = _StateRepDensePure(_np.zeros(state.data.shape, complex), state.state_space, state.basis)
        for f in self.factor_reps:
            output_state.data += f.acton(state).data
        return output_state

    def adjoint_acton(self, state):
        """ Act the adjoint of this operation matrix on an input state """
        output_state = _StateRepDensePure(_np.zeros(state.data.shape, complex), state.state_space, state.basis)
        for f in self.factor_reps:
            output_state.data += f.adjoint_acton(state).data
        return output_state

class OpRepEmbedded(OpRep):

    def __init__(self, state_space, target_labels, embedded_rep):

        state_space = _StateSpace.cast(state_space)
        iTensorProdBlks = [state_space.label_tensor_product_block_index(label) for label in target_labels]
        # index of tensor product block (of state space) a bit label is part of
        if len(set(iTensorProdBlks)) != 1:
            raise ValueError("All qubit labels of a multi-qubit operation must correspond to the"
                             " same tensor-product-block of the state space -- checked previously")  # pragma: no cover # noqa

        iTensorProdBlk = iTensorProdBlks[0]  # because they're all the same (tested above) - this is "active" block
        tensorProdBlkLabels = state_space.tensor_product_block_labels(iTensorProdBlk)
        # count possible *state-vector-space* indices of each component of the tensor product block
        numBasisEls = _np.array([state_space.label_udimension(l) for l in tensorProdBlkLabels], _np.int64)

        # Separate the components of the tensor product that are not operated on, i.e. that our
        # final map just acts as identity w.r.t.
        labelIndices = [tensorProdBlkLabels.index(label) for label in target_labels]
        actionInds = _np.array(labelIndices, _np.int64)
        assert(_np.product([numBasisEls[i] for i in actionInds]) == embedded_rep.dim), \
            "Embedded operation has dimension (%d) inconsistent with the given target labels (%s)" % (
                embedded_rep.dim, str(target_labels))

        #dim = state_space.udim
        nBlocks = state_space.num_tensor_product_blocks
        iActiveBlock = iTensorProdBlk
        nComponents = len(state_space.tensor_product_block_labels(iActiveBlock))
        embeddedDim = embedded_rep.dim  # a *unitary* dim - see .dim property above
        blocksizes = _np.array([_np.product(state_space.tensor_product_block_udimensions(k))
                                for k in range(nBlocks)], _np.int64)

        self.target_labels = target_labels
        self.embedded_rep = embedded_rep
        self.num_basis_els = numBasisEls
        self.action_inds = actionInds
        self.blocksizes = blocksizes

        num_basis_els_noop_blankaction = self.num_basis_els.copy()
        for i in self.action_inds: num_basis_els_noop_blankaction[i] = 1
        self.basisInds_noop_blankaction = [list(range(n)) for n in num_basis_els_noop_blankaction]

        # multipliers to go from per-label indices to tensor-product-block index
        # e.g. if map(len,basisInds) == [1,4,4] then multipliers == [ 16 4 1 ]
        self.multipliers = _np.array(_np.flipud(_np.cumprod([1] + list(
            reversed(list(self.num_basis_els[1:]))))), _np.int64)
        self.basisInds_action = [list(range(self.num_basis_els[i])) for i in self.action_inds]

        self.embeddedDim = embeddedDim
        self.ncomponents = nComponents  # number of components in "active" block
        self.active_block_index = iActiveBlock
        self.nblocks = nBlocks
        self.offset = sum(blocksizes[0:self.active_block_index])
        super(OpRepEmbedded, self).__init__(state_space)

    def _acton_other_blocks_trivially(self, output_state, state):
        offset = 0
        for iBlk, blockSize in enumerate(self.blocksizes):
            if iBlk != self.active_block_index:
                output_state.data[offset:offset + blockSize] = state.data[offset:offset + blockSize]  # identity op
            offset += blockSize

    def acton(self, state):
        output_state = _StateRepDensePure(_np.zeros(state.data.shape, complex), state.state_space, state.basis)
        offset = self.offset  # if rel_to_block else self.offset (rel_to_block == False here)

        for b in _itertools.product(*self.basisInds_noop_blankaction):  # zeros in all action-index locations
            vec_index_noop = _np.dot(self.multipliers, tuple(b))
            inds = []
            for op_b in _itertools.product(*self.basisInds_action):
                vec_index = vec_index_noop
                for i, bInd in zip(self.action_inds, op_b):
                    #b[i] = bInd #don't need to do this; just update vec_index:
                    vec_index += self.multipliers[i] * bInd
                inds.append(offset + vec_index)
            embedded_instate = _StateRepDensePure(state.data[inds],
                                                  state.state_space.create_subspace(self.target_labels), basis=None)
            embedded_outstate = self.embedded_rep.acton(embedded_instate)
            output_state.data[inds] += embedded_outstate.data

        #act on other blocks trivially:
        self._acton_other_blocks_trivially(output_state, state)
        return output_state

    def adjoint_acton(self, state):
        """ Act the adjoint of this gate map on an input state """
        #NOTE: Same as acton except uses 'adjoint_acton(...)' below
        output_state = _StateRepDensePure(_np.zeros(state.data.shape, complex), state.state_space, state.basis)
        offset = self.offset  # if rel_to_block else self.offset (rel_to_block == False here)

        for b in _itertools.product(*self.basisInds_noop_blankaction):  # zeros in all action-index locations
            vec_index_noop = _np.dot(self.multipliers, tuple(b))
            inds = []
            for op_b in _itertools.product(*self.basisInds_action):
                vec_index = vec_index_noop
                for i, bInd in zip(self.action_inds, op_b):
                    #b[i] = bInd #don't need to do this; just update vec_index:
                    vec_index += self.multipliers[i] * bInd
                inds.append(offset + vec_index)
            embedded_instate = _StateRepDensePure(state.data[inds],
                                                  state.state_space.create_subspace(self.target_labels), basis=None)
            embedded_outstate = self.embedded_rep.adjoint_acton(embedded_instate)
            output_state.data[inds] += embedded_outstate.data

        #act on other blocks trivially:
        self._acton_other_blocks_trivially(output_state, state)
        return output_state

class OpRepExpErrorgen(OpRep):

    def __init__(self, errorgen_rep):
        state_space = errorgen_rep.state_space
        self.errorgen_rep = errorgen_rep
        super(OpRepExpErrorgen, self).__init__(state_space)

    def errgenrep_has_changed(self, onenorm_upperbound):
        pass

    def acton(self, state):
        raise AttributeError("Cannot currently act with statevec.OpRepExpErrorgen - for terms only!")

    def adjoint_acton(self, state):
        raise AttributeError("Cannot currently act with statevec.OpRepExpErrorgen - for terms only!")

class OpRepRepeated(OpRep):
    def __init__(self, rep_to_repeat, num_repetitions, state_space):
        state_space = _StateSpace.cast(state_space)
        self.repeated_rep = rep_to_repeat
        self.num_repetitions = num_repetitions
        super(OpRepRepeated, self).__init__(state_space)

    def acton(self, state):
        """ Act this gate map on an input state """
        for i in range(self.num_repetitions):
            state = self.repeated_rep.acton(state)
        return state

    def adjoint_acton(self, state):
        """ Act the adjoint of this operation matrix on an input state """
        for i in range(self.num_repetitions):
            state = self.repeated_rep.adjoint_acton(state)
        return state

class OpRepLindbladErrorgen(OpRep):
    def __init__(self, lindblad_coefficient_blocks, state_space):
        super(OpRepLindbladErrorgen, self).__init__(state_space)
        self.Lterms = None
        self.Lterm_coeffs = None
        self.lindblad_coefficient_blocks = lindblad_coefficient_blocks