# Implementation of the CS-BuFLO defense (Cai et al., 14).
# Author: Giovanni Cherubin (@gchers)

import os
import sys
import math
import random
from collections import deque
from shutil import copyfile

import click
import numpy as np

from tqdm import tqdm
from numpy import median

# Direction
IN = -1
OUT = 1
MTU = 1
INF = float("inf")

# Upstream: from the current endpoint to the defence
# Downstream: from the other endpoint to the defence
# Define MODES of packets
UPSTREAM_APPLICATION_DATA = "upstream_application_data"
DOWNSTREAM_APPLICATION_DATA = "downstream_application_data"
ON_LOAD = "onLoad"
TIMEOUT = "timeout"

# Time is specified in ms
# NOTE: a QUIET_TIME too small causes the defence to fail,
# raising a TimeTravel exception: one endpoint was silent
# for too long, and then wants to catch up with the padding
# (junk) messages it didn't send.
# NOTE: The above problem was solved by using endpoint's
# "more_data" attribute, which is set to False only when
# the endpoint has no more data to transmit
QUIET_TIME = 2000
INIT_RHO = 200
MIN_RHO = 2 ** (-4) * 1000
MAX_RHO = 2 ** 3 * 1000

class TimeTravel(Exception):
    """This exception is raised when there are time inconsistencies.
    """

    pass

class Queue(deque, object):
    """A deque that allows seeing the element at the
    left of the queue without popping it.
    """

    def seeleft(self):
        """Returns the next element at the left of the
        queue, without removing it from the queue.

        Returns None if the queue is empty.
        """
        if self:
            return self[0]
        else:
            return None

class CSBuFLOEndpoint:
    """Implements one endpoint (client, server) of the CS-BuFLO defence as
    described in:
        http://pub.cs.sunysb.edu/~rob/papers/csbuflo.pdf
    """

    def __init__(self, output_trace, direction, initial_rho=0.2, quiet_time=QUIET_TIME):
        # Direction is OUT for client, IN for server
        self.direction = direction
        # Defended packets
        self.defended_trace = output_trace
        # Internal buffer for packets to send
        self.output_buff = Queue()  # Packets to send
        # Packets that have been sent
        # NOTE: don't rely on this to see which packets were sent,
        # because the function defend() pops from this queue.
        self.sent = Queue()
        # Bytes of real and padding (junk) traffic sent
        self.real_bytes = 0
        self.junk_bytes = 0
        self.last_site_response_time = None
        # onLoad and padding_done events
        self.on_load = False
        self.padding_done = False
        # The timeout it sampled uniformly in [0, 2*self.rho]
        # rho_stats are used to adjust rho
        self.initial_rho = initial_rho
        self.rho = initial_rho
        self.rho_stats = []
        self.timeout = initial_rho
        # "As a backup mechanism, the CS-BuFLO server considers
        # the website idle if quiet-time seconds pass without receiving
        # new data from the website. We used a quiet-time of 2
        # seconds in our prototype implementation."
        # (Section 4.5)
        self.quiet_time = quiet_time
        # Triggered to True when self.done_xmitting() == True
        self.done = False
        # This was not present in the original paper. Is is
        # set to False when no more data needs to be transmitted
        # from this endpoint
        self.more_data = True

    def next_timeout(self):
        ## No timeout if it finished transmitting
        # if self.done_xmitting():
        if self.padding_done:
            return INF
        # NOTE: need to check last sent packet time on defended_trace
        # rather than self.sent, as self.sent may be changed externally
        last_time = 0.0
        if self.defended_trace:
            for i in range(len(self.defended_trace) - 1, -1, -1):
                if self.defended_trace[i][1] == MTU * self.direction:
                    # last_time, _ = self.defended_trace[i]
                    last_time, _ = self.defended_trace[i]
                    break

        return last_time + self.timeout

    def process(self, packet, mode):
        # print('process {} {}'.format(packet, mode))
        # print('rho {}'.format(self.rho))
        # print('timeout {}'.format(self.timeout))
        time, size = packet

        if mode == UPSTREAM_APPLICATION_DATA:
            # this endpoint is trying to send data
            self.output_buff.append((time, size))
            self.real_bytes += abs(size)
            self.last_site_response_time = time
            # self.padding_done = False
        elif mode == DOWNSTREAM_APPLICATION_DATA:
            # would pass data back to this endpoint
            self.rho_stats.append(None)
            self.on_load = False
            # NOTE: This is only done for the server
            # endpoint (i.e., direction == IN)
            # This was not specified by the original paper
            # but without this we would encounter instable
            # situations
            if self.direction == IN:
                self.padding_done = False
        elif mode == ON_LOAD:
            self.on_load = True
        elif mode == TIMEOUT:
            if self.output_buff:
                self.rho_stats.append(time)
            padding = self.cs_send(time)
            self.junk_bytes += padding

        if self.done_xmitting(time):
            self.padding_done = True
        else:
            # Set rho to the average time between sends to client
            if self.rho == INF:
                self.rho = self.initial_rho
            elif self.crossed_threshold():
                self.rho = self.rho_estimator()
                self.rho_stats = []

            if mode == TIMEOUT:
                # Random in [0, 2*rho]
                self.timeout = random.random() * 2 * self.rho

    def cs_send(self, time):
        padding = 0

        if self.output_buff:
            _, size = self.output_buff.popleft()
            # We use the timeout time
            packet = (time, size)
        else:
            # Padding packet
            packet = (time, self.direction * MTU)
            padding = 1

        self.sent.append(packet)
        self.defended_trace.append(packet)

        return padding

    def padding_finished(self):
        """Using "payload padding", which pads until the bytes
        sent are a multiple of the real bytes.

        Implements "Payload padding" as in paper.
        """
        if self.real_bytes == 0:
            return False

        total_bytes = self.real_bytes + self.junk_bytes
        next_multiple = 2 ** math.ceil(math.log(self.real_bytes, 2))

        return total_bytes % next_multiple == 0

    def done_xmitting(self, cur_time):
        condition1 = self.padding_finished() or self.crossed_threshold()
        # NOTE: The paper writes "length(output_buff) <- 0" as first
        # condition. Since the expression makes no sense, I assume
        # the authors meant "length(output_buff) == 0".
        condition2 = not self.output_buff
        condition3 = self.channel_idle(cur_time)
        # NOTE: I add the following condition, because if one
        # endpoint gets silent for too long and then wants to transmit
        # again the defence would fail
        condition4 = not self.more_data

        # print(self.direction, condition1, condition2, condition3, condition4)

        return condition1 and condition2 and condition3 and condition4

    def channel_idle(self, cur_time):
        if self.last_site_response_time is not None:
            is_quiet = self.last_site_response_time + self.quiet_time < cur_time
        else:
            is_quiet = False

        return self.on_load or is_quiet

    def rho_estimator(self):
        it = []
        for i in range(len(self.rho_stats) - 1):
            if self.rho_stats[i] is not None and self.rho_stats[i + 1] is not None:
                it.append(self.rho_stats[i + 1] - self.rho_stats[i])

        if not it:
            return self.rho

        med = median(it)
        if not med:
            rho = self.rho

        rho = 2 ** math.floor(math.log(med, 2))

        # NOTE: to my understanding, the original CS-BuFLO
        # implemented in SSH has a minimum and maximum RHO
        # values. By doing some experiments I observed the
        # following values allow containing overheads
        if rho < MIN_RHO:
            rho = MIN_RHO
        elif rho > MAX_RHO:
            rho = MAX_RHO

        return rho

    def crossed_threshold(self):
        # From Algorithm 1 in the paper, this function seems to be
        # called on two arguments: (real_bytes, junk_bytes).
        # However, the function only accepts one argument (Algorithm 4).
        # Thankfully, DONE_XMITTING() in Algorithm 4 calls the function
        # on real_bytes + junk_bytes, so I assume this is what was
        # intended in Algorithm 1.
        x = self.real_bytes + self.junk_bytes

        # NOTE: I'm adding this, because otherwise (_by design_)
        # when x == MTU (which does happen) the log is indefinite
        if x == MTU or x == 0:
            return False

        # The following comment does not apply anymore, and I'll
        # only keep it for record.
        ## NOTE: This is an edit from the original function.
        ## As in Figure 1, it returns True only if the number
        ## of bytes sent so far is a multiple of 2^k, for
        ## some integer k
        ##log = math.log(x, 2)
        ##return log == math.floor(log)

        return math.floor(math.log(x - MTU, 2)) < math.floor(math.log(x, 2))

class CSBuFLO:
    def __init__(self, initial_rho=INIT_RHO):
        self.initial_rho = initial_rho

    def reset(self):
        self.defended_trace = Queue()
        self.client = CSBuFLOEndpoint(self.defended_trace, OUT, self.initial_rho)
        self.server = CSBuFLOEndpoint(self.defended_trace, IN, self.initial_rho)
        # For each endpoint (client, server), there are two
        # queues of packets: those read from the trace and those
        # that were actually sent by the defence (output).
        # The output ones can be found as client.sent, server.sent,
        # the others we define as follows
        self.client_packets = Queue()
        self.server_packets = Queue()

    def defend(self, packets):
        self.reset()

        # Add packets to the respective pipelines
        for t, s in packets:
            if s > 0:
                self.client_packets.append((t, s))
            else:
                self.server_packets.append((t, s))

        # Record what events happen for debugging purposes
        nevents = {}

        # Keep time, to check for errors
        self.cur_time = 0.0

        running = True
        while running:
            # Detect onLoad
            cond1 = not self.client_packets and not self.client.output_buff
            cond2 = not self.server_packets and not self.server.output_buff
            if cond1 and cond2:
                # print('onLoad')
                self.server.process((self.cur_time, None), ON_LOAD)
                self.client.process((self.cur_time, None), ON_LOAD)

            # Notify endpoint if no more data to transmit
            if not self.client_packets:
                self.client.more_data = False
            if not self.server_packets:
                self.server.more_data = False

            # We break if either there's no more data to transmit
            # or self.process_next() returns None; the latter happens
            # when there are no more packets to transmit and
            # padding is done.
            if self.client.done_xmitting(self.cur_time) and self.server.done_xmitting(
                self.cur_time
            ):
                break

            # Keep log of events
            next_event = self.process_next()
            if next_event not in nevents:
                nevents[next_event] = 1
            else:
                nevents[next_event] += 1

            # print('client out: {}'.format(self.client.output_buff))
            # print('server out: {}'.format(self.server.output_buff))
            ##print('defended: {}'.format(defended_trace))
            # print('client rho: {}'.format(self.client.timeout))
            # print('server rho: {}'.format(self.server.timeout))

        return list(self.defended_trace)

    def process_next(self):
        """Process next event.

        Get the time of each event, and process the one
        that comes first.
        """
        # Find what is the next event
        time_c_read_packet = self.client_packets.seeleft()
        time_s_read_packet = self.server_packets.seeleft()
        # Little hack to use the "min" expression later on both
        # packets and timeouts
        c_timeout = (self.client.next_timeout(),)
        s_timeout = (self.server.next_timeout(),)

        t_events = [time_c_read_packet, time_s_read_packet, c_timeout, s_timeout]

        next_packet = min([x for x in t_events if x is not None], key=lambda x: x[0])
        next_event = t_events.index(next_packet)

        # DEBUG
        # print('t-events {}'.format(t_events))
        # print('next {}'.format(next_event))
        # print(self.client.timeout)
        # print(len(self.client_packets))

        # Check the time is consistent
        next_time = next_packet[0]
        # NOTE: if next_time is INF, it means both endpoints
        # have been silent for a while. Try increasing QUIET_TIME.
        if next_time < self.cur_time:
            print("I travelled back in time:")
            print("{} => {}".format(self.cur_time, next_time))
            print(t_events)
            raise TimeTravel
        self.cur_time = next_time

        # Process the event
        process_event = {
            0: self.client_read,
            1: self.server_read,
            2: self.client_timeout,
            3: self.server_timeout,
        }

        process_event[next_event]()
        # After event, send packet (if any) to the other
        # endpoint
        if self.client.sent.seeleft() is not None:
            self.client_sent()
        if self.server.sent.seeleft() is not None:
            self.server_sent()

        return next_event

    def client_read(self):
        packet = self.client_packets.popleft()
        self.client.process(packet, UPSTREAM_APPLICATION_DATA)

    def server_read(self):
        packet = self.server_packets.popleft()
        self.server.process(packet, UPSTREAM_APPLICATION_DATA)

    def client_sent(self):
        packet = self.client.sent.popleft()
        self.server.process(packet, DOWNSTREAM_APPLICATION_DATA)

    def server_sent(self):
        packet = self.server.sent.popleft()
        self.client.process(packet, DOWNSTREAM_APPLICATION_DATA)

    def client_timeout(self):
        packet = (self.client.next_timeout(), None)
        self.client.process(packet, TIMEOUT)

    def server_timeout(self):
        packet = (self.server.next_timeout(), None)
        self.server.process(packet, TIMEOUT)

def normalise_timings(packets):
    """Returns a packet sequence with sorted packets, and
    timing starting at 0.
    """
    if not packets:
        return []
    res = sorted(packets, key=lambda x: x[0])
    min_time = res[0][0]

    for i in range(len(res)):
        t, s = res[i]
        res[i] = (t - min_time, s)
        if t == INF:
            print("Try increasing QUIET_TIME")
            raise TimeTravel

    return res

@click.command()
@click.option("--data_path")
@click.option("--out_path")
def main(data_path, out_path):
    DEFENDED = out_path
    dataset = data_path
    outdirectory = DEFENDED

    if not os.path.exists(outdirectory):
        os.makedirs(outdirectory)

    unmod = []
    mod = []
    added = []
    count = 0

    # Defend
    for fname in tqdm(os.listdir(dataset)):
        infname = os.path.join(dataset, fname)
        outfname = os.path.join(outdirectory, fname)
        # Skip open world traces
        if "-" not in fname:
            # copyfile(infname, outfname)
            continue

        packets = []
        with open(infname, "r") as f:
            for x in f.readlines():
                t, s = x.split("\t")
                t = float(t) * 1000.0  # To milliseconds
                s = int(s)
                packets.append((t, s))
        success = False
        while not success:
            try:
                defended = CSBuFLO().defend(packets)
                success = True
            except TimeTravel:
                pass

        defended = normalise_timings(defended)

        # Store defended trace
        with open(outfname, "w") as f:
            for t, s in defended:
                t /= 1000.0  # To seconds
                s = 1 if s > 0 else -1
                f.write(repr(t) + "\t" + repr(s) + "\n")

if __name__ == "__main__":
    main()


'''
Essential functions and classes for calibration. See, in particular:

- `OPT_BOUNDS`, which has lower and upper bounds on calibration parameters

**You must create a configuration JSON file before calibrating L4C.** There
is a template available in the directory:

    pyl4c/data/fixtures/files

The optimization routines here are best used with multiple `--trials`, i.e.,
multiple, random initial parameter values. This can help avoiding falling into
a local optimum when a better solution set is available. However, if you are
repeatedly seeing the message:

    Error in objective function; restarting...

Most likely, the bounds on your parameters are producing initial parameter
values that are totally unreasonable.

Note that calibration of L4C (i.e., updating the BPLUT) can be performed using
the command-line interface in `pyl4c.apps.calibration.main`, for example:

    # Build the scratch dataset needed for calibration
    python main.py setup

    # You can get a preview of what filtering the data would look like...
    python main.py pft <pft_number> filter-preview gpp <window_size>
    python main.py pft <pft_number> filter-preview reco <window_size>

    # Optionally, filter the tower datasets to remove spurious spikes
    python main.py filter-all gpp <window_size>
    python main.py filter-all reco <window_size>

    # Run the GPP calibration for a given Plant Functional Type (PFT)
    python main.py pft <pft_number> tune-gpp

    # Run the RECO calibration for a given Plant Functional Type (PFT)
    python main.py pft <pft_number> tune-reco

    # Finally, to dump the updated BPLUT into a CSV or Python pickle file...
    python main.py bplut pickle <output_path> --version-id=<version_id>

    # Get help on any command with --help, for example:
    python main.py plot-gpp --help
'''

import os
import pickle
import nlopt
import netCDF4 # Necessary to import this before HDF5 due to a bug
import h5py
import numpy as np
from collections import OrderedDict
from scipy import optimize
from pyl4c import suppress_warnings
from pyl4c.stats import detrend, rmsd, sum_of_squares
from pyl4c.science import k_mult

# Constrained optimization bounds
OPT_BOUNDS = {
    'gpp': ( # lue, tmin0, tmin1, vpd0, vpd1, smrz0, smrz1, ft0
        np.array((0.5, 230, 276,    0,  1501,  0,  30.1, 0.)), # Lower bound
        np.array((4.0, 275, 320, 1500, 10000, 30, 100,   1.))), # Upper bound
    'reco': ( # CUE, beta_tsoil, smsf0, smsf1
        np.array((0.0,   1,    0,  25)),
        np.array((0.7, 800, 24.9, 100)))
}

class BPLUT(object):
    '''
    Represents a Biome Properties Look-Up Table (BPLUT) with PFT classes along
    the rows and parameters along the columns.

    If initialized with a `params_dict`, these are the values of the BPLUT.
    If initialized with an `hdf5_path` but without a `params_dict`, the
    parameters are read-in from the HDF5 file.
    '''
    _labels = [
        'LUE', 'CUE', 'tmin0', 'tmin1', 'vpd0', 'vpd1', 'smrz0', 'smrz1',
        'smsf0', 'smsf1', 'ft0', 'ft1', 'tsoil', 'decay_rates0',
        'decay_rates1', 'decay_rates2', 'f_metabolic', 'f_structural'
    ]
    _npft = 10 # Number of (possible) PFT classes
    _valid_pft = range(1, 9) # Canonical range of valid PFTs

    def __init__(
            self, params_dict = None, labels = None, hdf5_path = None,
            hdf5_group = 'BPLUT'):
        '''
        Parameters
        ----------
        params_dict : dict
            A BPLUT to initialize the new BPLUT
        labels : tuple or list
            Names of the parameters
        hdf5_path : str
            Path to an HDF5 file to use as a temporary store
        hdf5_group : str
            Field name with which to store data in HDF5 file
            (Default: "BPLUT")
        '''
        self.hdf5_group = hdf5_group
        self.hdf5_path = hdf5_path
        if labels is not None:
            print('WARNING: Parameter names ending with a number are assumed to be bounds on ramp functions!')
            self._labels = labels
        # Create an in-memory parameter dictionary
        empty = self._empty_dict(self.labels)
        if params_dict is None:
            init_data = empty # No prior BPLUT, use empty table
        else:
            init_data = params_dict.copy()
            # IMPORTANT: Make sure prior BPLUT has all the necessary params
            for key in empty.keys():
                init_data.setdefault(key, empty[key])
        # Optionally, maintain a file BPLUT
        if hdf5_path is not None:
            # Restore from the file data, filling in NaNs with initial
            if os.path.exists(hdf5_path):
                with h5py.File(hdf5_path, 'r') as hdf:
                    if hdf5_group in hdf.keys():
                        init_data = self.hdf5_restore(hdf, init_data = init_data)
            # Then, if file dataset doesn't exist, create a new one and store
            #   the initial data
            with h5py.File(hdf5_path, 'a') as hdf:
                self.hdf5_flush(hdf, data = init_data)
        self.data = init_data

    @property
    def data(self):
        'The parameters dictionary or `dict` instance'
        return self._data

    @data.setter
    def data(self, data):
        self._data = data

    @property
    def labels(self):
        'Names of the free parameters'
        return self._labels

    def __getitem__(self, key):
        return self.data[key]

    def __setitem__(self, key, value):
        self.data[key] = value

    def _empty_dict(self, labels = None, dtype = np.float32):
        '''
        Given sequence of labels, convert to legacy, human-readable dict,
        e.g.:

            { 'LUE': array([[ nan, 1.17, ..., nan ]]), ... }
        '''
        labels_dedupe = self._canonical(labels)
        result = dict()
        for name in labels_dedupe:
            size = len(list(filter(lambda x: x.startswith(name), labels)))
            result[name] = np.ones((size, self._npft), dtype = dtype) * np.nan
        return result

    def _canonical(self, labels = None):
        '''
        Returns a short list of labels, without duplicates. Specifically,
        a list of labels like, e.g., ("tmin0", "tmin1", "decay_rates1")
        becomes ("tmin", "decay_rates").

        Parameters
        ----------
        labels : tuple or list or None

        Returns
        -------
        list
        '''
        labels = self.labels if labels is None else labels
        # Remove numeric suffixes and de-duplicate the list
        return list(OrderedDict([
            (p.strip('0123456789'), 0) for p in labels
        ]).keys())

    def flat(self, pft, labels = None):
        '''
        Retrieves a flat list of parameters for a specific PFT.

        Parameters
        ----------
        pft : int
            Numeric code of the PFT for which to return parameter values
        labels : tuple or list
            (Optional) A sequence of parameter names desired, if not all;
            defaults to returning all parameters

        Returns
        -------
        numpy.ndarray
        '''
        labels_dedupe = self._canonical(labels)
        return np.hstack([
            self.data[p][:,pft].ravel() if p in self.data.keys() else np.nan
            for p in labels_dedupe
        ])

    def hdf5_flush(self, hdf, data = None):
        '''
        Writes the current BPLUT to an HDF5 file.

        Parameters
        ----------
        hdf : h5py.File
            HDF5 file open for writing
        data : dict
        '''
        assert hdf.mode != 'r', 'File not open for writing!'
        data = self.data if data is None else data
        if self.hdf5_group not in hdf.keys():
            hdf.create_group(self.hdf5_group)
        for key, value in data.items():
            if key.startswith('_'):
                continue # Skip keys that are not parameter names
            field = '%s/%s' % (self.hdf5_group, key)
            if key not in hdf[self.hdf5_group].keys():
                hdf.create_dataset(field, value.shape, np.float32, value)
            else:
                # Overwrite NaNs in the file data
                _value = hdf[field][:]
                hdf[field][:] = np.where(np.isnan(_value), value, _value)
        hdf.flush()

    def hdf5_restore(self, hdf, init_data = None, dtype = np.float32):
        '''
        Reads in the BPLUT table stored in the HDF5 file.

        Parameters
        ----------
        hdf : h5py.File
        init_data : dict
            Initital data; will be over-written by HDF5 file contents

        Returns
        -------
        dict
        '''
        data = dict() if init_data is None else init_data
        for key in hdf[self.hdf5_group].keys():
            # Update the in-memory BPLUT
            from_hdf5 = hdf[self.hdf5_group][key][:]
            data[key] = np.where(
                ~np.isnan(from_hdf5), from_hdf5, init_data.get(key)
            ).astype(dtype)
        return data

    def pickle(self, output_path, version_id = None):
        '''
        Writes the current BPLUT parameters, as a dictionary, to a pickle
        file.

        Parameters
        ----------
        output_path : str
            The output path for the pickle file (*.pickle)
        version_id : str
            (Optional) The version identifier for this BPLUT
        '''
        with open(output_path, 'wb') as file:
            output = self.data.copy()
            if version_id is not None:
                output['_version'] = version_id
            pickle.dump(output, file)

    def show(self, pft, param, precision = 2):
        '''
        Prints the current BPLUT parameters for a given PFT, the values of a
        given parameter for all PFTs, or the value of a specific PFT-parameter
        combination.

        Parameters
        ----------
        pft : int or None
            The PFT class
        param : str or None
            The name of the parameter
        precision : int
            Decimal precision to use for printing numbers
        '''
        assert not (pft is None and param is None),\
            'Either one or both must be specified: --pft or --param'
        set_of_labels = self._canonical() if param is None else [param]
        set_of_pfts = self._valid_pft if pft is None else [pft]
        for each in set_of_labels:
            assert each in self._canonical(), 'Unrecognized parameter: %s' % each
        for each in set_of_pfts:
            assert each in range(0, self._npft), 'PFT code out of range'
        pad = max(len(l) for l in set_of_labels) + 2
        fmt_string = '{:>%d} {:>%d}' % (pad, 5 + precision)
        for pft in set_of_pfts:
            print('BPLUT parameters for PFT %d:' % pft)
            for label in set_of_labels:
                param_values = self.data[label][:,pft]
                for i, value in enumerate(param_values):
                    # If there are multiple values for a parameter (group),
                    #   append a number to the end of the label
                    if len(param_values) > 1:
                        prefix = '%s%d:' % (label, i)
                    else:
                        prefix = '%s:' % label
                    print(
                        fmt_string.format(prefix, ('%%.%df' % precision) % value))

    def update(self, pft, values, labels, flush = True):
        '''
        Updates the BPLUT with the specified parameters for a single PFT.

        Parameters
        ----------
        pft : int
            The PFT class
        values : tuple or list
            Sequence of parameter values, one for each parameter named in
            `labels`
        labels : tuple or list
            Sequence of parameter names, one for each value in `values`
        flush : bool
            True to write the result to disk (attached HDF5 file storage)
            (Default: True)
        '''
        assert len(values) == len(labels),\
            'Vectors of values and parameter labels must have the same length'
        if flush:
            assert self.hdf5_path is not None,\
                'No HDF5 file storage is attached'
            hdf = h5py.File(self.hdf5_path, 'a')
        for i, name in enumerate(labels):
            abbrv = name.strip('0123456789')
            # In case parameter has multiple levels, like a ramp function
            #   (e.g., "smsf0" and "smsf1" are two rows)
            dupes = list(filter(lambda x: x.startswith(abbrv), self.labels))
            j = dupes.index(name) # If it has one level (e.g., "LUE"), j = 0
            self.data[abbrv][j,pft] = values[i]
            if flush:
                path = '%s/%s' % (self.hdf5_group, abbrv)
                hdf[path][:,pft] = self.data[abbrv][:,pft]
        if flush:
            hdf.flush()
            hdf.close()

class ModelParameters(OrderedDict):
    '''
    Convenience wrapper for an OrderedDict, allowing both vectorized and
    keyword access to model parameters.

    Parameters
    ----------
    group : str
        Name of this model parameters group, usually the name of the model or
        sub-model to which they belong
    *params : spotpy.parameter
        One or more parameters
    '''
    def __init__(self, group, *params):
        self._group = group
        # Create {name: spotpy.parameter, ...} dictionary
        super().__init__(**dict([(p.name, p) for p in params]))

class GenericOptimization(object):
    '''
    A more generic and expansive tool for optimization; includes many more
    algorithms for minimization/ maximization problems, including sequential
    quadratic programming (SQP), which is the default here and is closest to
    what is performed in Matlab's `fmincon`. Despite the similarity to `fmincon`,
    SQP will tend to deviate strongly from the initial parameters derived via
    fmincon. This solver is SLOW for gradient descent methods relative to
    `scipy.optimize.least_squares()`, because the gradient is calculated with
    a finite element approach.

        opt = GenericOptimization(residuals, OPT_BOUNDS['gpp'],
            step_size = (0.01, 0.1, 0.1, 1, 1, 0.1, 0.1, 0.05))
        opt.solve(init_params)

    See: https://nlopt.readthedocs.io/en/latest/NLopt_Python_Reference/

    Parameters
    ----------
    func : function
        Function to calculate the residuals
    bounds : list or tuple
        2-element sequence of (lower, upper) bounds where each element is an
        array
    method : str
        One of the nlopt algorithms
    step_size : list or tuple or numpy.ndarray
        Sequence of steps to take in gradient descent; not needed for
        derivative-free methods
    verbose : bool
        True to print all output to the screen
    '''
    def __init__(
            self, func, bounds, method: int = nlopt.LD_SLSQP,
            step_size = None, verbose = True):
        # https://nlopt.readthedocs.io/en/latest/NLopt_Algorithms/#slsqp
        assert isinstance(method, int), 'Did not recognize "method" argument'
        self._bounds = bounds
        self._method = method
        self._residuals = func
        self._step_size = step_size
        self._verbose = verbose

    def solve(self, init_params, ftol = 1e-8, xtol = 1e-8, maxeval = 500):
        '''
        Using the sum-of-squared errors (SSE) as the objective function,
        solves a minimization problem.

        Parameters
        ----------
        init_params : list or tuple or numpy.ndarray
            Sequence of starting parameters (or "initial guesses")
        ftol : float
        xtol : float
        maxeval : int
            Maximum number of objective function evaluations

        Returns
        -------
        numpy.ndarray
        '''
        @suppress_warnings
        def sse(x):
            return np.power(self._residuals(x), 2).sum()

        @suppress_warnings
        def objf(x, grad):
            if grad.size > 0:
                # Approximate the gradient using finite element method
                grad[...] = optimize.approx_fprime(
                    x, sse, self._step_size)
            return sse(x)

        opt = nlopt.opt(self._method, len(init_params))
        # https://nlopt.readthedocs.io/en/latest/NLopt_Python_Reference/#localsubsidiary-optimization-algorithm
        if self._method == nlopt.G_MLSL_LDS:
            opt.set_local_optimizer(
                nlopt.opt(nlopt.LN_COBYLA, len(init_params)))
        opt.set_min_objective(objf)
        opt.set_lower_bounds(self._bounds[0])
        opt.set_upper_bounds(self._bounds[1])
        opt.set_ftol_abs(ftol)
        opt.set_xtol_abs(xtol)
        opt.set_maxeval(maxeval)
        if self._verbose:
            print('Solving...')
        return opt.optimize(init_params)

def cbar(rh, k_mult, q_rh = 75, q_k = 50):
    '''
    Calculates "Cbar," the time-constant upper quantile of the RH/Kmult
    ratio. Where Kmult is >/= `q_k`, return the `q_rh` quantile of RH/Kmult;
    intended for T x N arrays where T is the number of time steps and
    N is the number of (flux tower) sites.

    Parameters
    ----------
    rh : numpy.ndarray
        (T x N) vector of heterotrophic respiration
    k_mult : numpy.ndarray
        (T x N) vector of Kmult
    q_rh : float
        Percentile of RH/Kmult to return
    q_k : float
        Percentile of Kmult below which RH/Kmult values are masked

    Returns
    -------
    numpy.float64
    '''
    cutoff = np.apply_along_axis(
        np.percentile, 0, k_mult, q = q_k).reshape((1, k_mult.shape[1]))
    return np.nanpercentile(
        np.where(k_mult >= cutoff,
            np.divide(rh, np.where(k_mult == 0, np.nan, k_mult)), np.nan),
        q = q_rh, axis = 0)

def reco(params, tsoil, smsf, reco_tower, gpp_tower, q_rh = 75, q_k = 50):
    '''
    Calculate empirical ecosystem respiration, RECO, based on current model
    parameters and the inferred soil organic carbon (SOC) storage; i.e., this
    calculation should be used in model calibration when SOC is not a priori
    known, see `pyl4c.apps.calibration.cbar()`. The expected model parameter
    names are "CUE" for the carbon use efficiency of plants.

    Parameters
    ----------
    params : dict
        A dict-like data structure with named model parameters
    tsoil : numpy.ndarray
        (T x N) vector of soil temperature (deg K), where T is the number of
        time steps, N the number of sites
    smsf : numpy.ndarray
        (T x N) vector of surface soil wetness (%), where T is the number of
        time steps, N the number of sites
    reco_tower : numpy.ndarray
        (T x N) vector of observed RECO from eddy covariance tower sites
    gpp_tower : numpy.ndarray
        (T x N) vector of observed GPP from eddy covariance tower sites
    q_rh : int
        The percentile of RH/Kmult to use in calculating Cbar
    q_k : int
        The percentile of Kmult below which RH/Kmult values are masked

    Returns
    -------
    numpy.ndarray
    '''
    # Calculate RH as (RECO - RA) or (RECO - (faut * GPP));
    #   globals "reco_tower", "gpp_tower"
    ra = ((1 - params['CUE']) * gpp_tower)
    rh = reco_tower - ra
    rh = np.where(rh < 0, 0, rh) # Mask out negative RH values
    # Compute Cbar with globals "q_rh" and "q_k"
    kmult0 = k_mult(params, tsoil, smsf)
    cbar0 = cbar(rh, kmult0, q_rh, q_k)
    return ra + (kmult0 * cbar0)

def report_fit_stats(obs, pred, weights = np.array([1]), verbose = True):
    '''
    Reports the RMSE, ubRMSE, and Bias for observed and predicted values.

    Parameters
    ----------
    obs : numpy.ndarray
        Vector of observed ("true") values
    pred : numpy.ndarray
        Vector of predicted values
    weights : numpy.ndarray
        (Optional) Vector of weights for each sample

    Returns
    -------
    tuple
        (R-squared, RMSE, ubRMSE, Bias)
    '''
    y = np.apply_along_axis(detrend, 0, obs, fill = True)
    yhat = np.apply_along_axis(detrend, 0, pred, fill = True)
    rmse = rmsd(obs, pred, weights = weights)
    ubrmse = rmsd(y, yhat, weights = weights)
    bias = np.nanmean(np.subtract(pred, obs))
    mask = np.logical_or(np.isnan(obs), np.isnan(pred))
    r_squared = 1 - np.divide(
        sum_of_squares(
            obs[~mask], pred[~mask], add_intercept = False, which = 'sse'),
        sum_of_squares(
            obs[~mask], pred[~mask], add_intercept = False, which = 'sst'))
    if verbose:
        print('Fit statistics:')
        print('--    R^2: %s' % ('%.3f' % r_squared).rjust(6))
        print('--   RMSE: %s' % ('%.3f' % rmse).rjust(6))
        print('-- ubRMSE: %s' % ('%.3f' % ubrmse).rjust(6))
        print('--   Bias: %s' % ('%.3f' % bias).rjust(6))
    return (r_squared, rmse, ubrmse, bias)

def solve_least_squares(func, init_params, labels, bounds, **kwargs):
    '''
    Apply constrained, non-linear least-squares optimization. Mostly a
    wrapper for `scipy.optimize.least_squares()`.

    Parameters
    ----------
    func : function
        Function to calculate the residuals
    init_params : list or tuple or numpy.ndarray
        Sequence of starting parameters (or "initial guesses")
    labels : list or tuple or numpy.ndarray
        Sequence of parameter names
    bounds : list or tuple
        2-element sequence of (lower, upper) bounds where each element is an
        array

    Returns
    -------
    scipy.optimize.OptimizeResult
    '''
    # Update the optimization settings; this loss function produces
    #   an estimate for the FT multiplier that is closest to prior
    kwargs.setdefault('loss', 'arctan')
    kwargs.setdefault('method', 'trf')
    kwargs.setdefault('max_nfev', 500)
    kwargs.setdefault('ftol', 1e-8)
    kwargs.setdefault('xtol', 1e-8)
    kwargs.setdefault('gtol', 1e-8)
    try:
        solution = optimize.least_squares(
            func, init_params, bounds = bounds, **kwargs)
    except ValueError:
        below = [
            labels[i]
            for i in np.argwhere(init_params < bounds[0]).flatten().tolist()
        ]
        above = [
            labels[i]
            for i in np.argwhere(init_params > bounds[1]).flatten().tolist()
        ]
        if np.isnan(init_params).any():
            raise ValueError(
                'Error in candidate parameter values; residual function probably returning NaNs')
        else:
            raise ValueError(
                '"Infeasibility" error; check lower bound on %s; upper bound on %s' % (
                '(None)' if len(below) == 0 else ', '.join(below),
                '(None)' if len(above) == 0 else ', '.join(above)))
    return solution


#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function

import io
import json
import math
import os
import shutil
import struct

import h5py
import numpy as np
import obspy
import xarray

# LASIF can already deal with the binary SES3D models. Thus we can utilize
# this here!
from lasif.ses3d_models import RawSES3DModelHandler
from lasif.scripts.lasif_cli import _find_project_comm

def binary_ses3d_to_hdf5_model(input_folder, lasif_project, output_filename):
    """
    Function converting a binary SES3D model consisting of many files to an
    HDF5 model.

    Requires access to a LASIF project that determines the potentially
    rotated geometry. Not super clean but workable.

    :param input_folder: The folder containing the input model.
    :param lasif_project: The folder with the LASIF project.
    :param output_filename: The output filename.
    """
    assert not os.path.exists(output_filename), \
        "'%s' already exists" % output_filename
    _dirname = os.path.dirname(output_filename)
    if not os.path.exists(_dirname) and _dirname:
        os.makedirs(_dirname)

    # We need the project to get the domain definition which stores the
    # rotation settings.
    comm = _find_project_comm(lasif_project, read_only_caches=True)

    if any(_i.startswith("grad_") for _i in os.listdir(input_folder)):
        model_type = "kernel"
    else:
        model_type = "earth_model"

    m = RawSES3DModelHandler(
        directory=input_folder, domain=comm.project.domain,
        model_type=model_type)

    f = h5py.File(output_filename)

    try:
        data_group = f.create_group("data")

        if model_type == "earth_model":
            # We will also store A, C, and Q which we don't invert for but
            # have to take into account in any case.
            components = ["vp", "vsh", "vsv", "rho", "A", "C"]
            # Q might not exist.
            if "Q" in m.components:
                components.append("Q")
        elif model_type == "kernel":
            components = ["grad_cp", "grad_csh", "grad_csv", "grad_rho"]
        else:
            raise NotImplementedError

        # Make it compatible with seismopt.
        rename_dict = {
            "grad_cp": "vp",
            "grad_csh": "vsh",
            "grad_csv": "vsv",
            "grad_rho": "rho",
        }

        for c in components:
            m.parse_component(c)
            _d = xarray.DataArray(
                np.require(m.parsed_components[c],
                           requirements=["C_CONTIGUOUS"]),
                coords=[90.0 - m.collocation_points_lats[::-1],
                        m.collocation_points_lngs,
                        (6371.0 - m.collocation_points_depth) * 1000.0],
                dims=["colatitude", "longitude", "radius_in_m"])

            # Write to HDF5 file.
            if c in rename_dict:
                c = rename_dict[c]

            data = _d.data

            # Make sure it is g\cm^3 in the hdf5 files.
            if c == "rho":
                if data.mean() > 1000.0:
                    data /= 1000.0

            data_group[c] = np.require(data, dtype=np.float32)

            data_group[c].attrs["variable_name"] = \
                np.string_((c + "\x00").encode())

        # Write coordinate axes.
        f["coordinate_0"] = np.require(_d.colatitude.data, dtype=np.float32)
        f["coordinate_0"].attrs["name"] = \
            np.string_(("colatitude" + "\x00").encode())
        f["coordinate_1"] = np.require(_d.longitude.data, dtype=np.float32)
        f["coordinate_1"].attrs["name"] = \
            np.string_(("longitude" + "\x00").encode())
        f["coordinate_2"] = np.require(_d.radius_in_m.data, dtype=np.float32)
        f["coordinate_2"].attrs["name"] = \
            np.string_(("radius_in_m" + "\x00").encode())

        # Create dimension scales.
        for c in components:
            if c in rename_dict:
                c = rename_dict[c]
            f["data"][c].dims[0].label = "colatitude"
            f["data"][c].dims[1].label = "longitude"
            f["data"][c].dims[2].label = "radius_in_m"

            f["data"][c].dims.create_scale(f["coordinate_0"], "values")
            f["data"][c].dims.create_scale(f["coordinate_1"], "values")
            f["data"][c].dims.create_scale(f["coordinate_2"], "values")

            f["data"][c].dims[0].attach_scale(f["coordinate_0"])
            f["data"][c].dims[1].attach_scale(f["coordinate_1"])
            f["data"][c].dims[2].attach_scale(f["coordinate_2"])

        # Also add some meta information.
        _meta = f.create_group("_meta")
        model_name = os.path.split(os.path.normpath(os.path.abspath(
            input_folder)))[-1]
        _meta.attrs["model_name"] = np.string_((model_name + "\x00").encode())

        # Everything needed to reconstruct the domain objects.
        _domain = _meta.create_group("domain")
        d = comm.project.domain
        _domain.attrs["min_longitude"] = d.min_longitude
        _domain.attrs["max_longitude"] = d.max_longitude
        _domain.attrs["min_latitude"] = d.min_latitude
        _domain.attrs["max_latitude"] = d.max_latitude
        _domain.attrs["min_depth_in_km"] = d.min_depth_in_km
        _domain.attrs["max_depth_in_km"] = d.max_depth_in_km
        _domain.attrs["rotation_axis"] = d.rotation_axis
        _domain.attrs["rotation_angle_in_degree"] = d.rotation_angle_in_degree
        _domain.attrs["boundary_width_in_degree"] = d.boundary_width_in_degree

        # We also need to store the boxfile.
        _meta.create_dataset("boxfile",
                             data=np.fromfile(m.boxfile, dtype=np.uint8))
    finally:
        try:
            f.close()
        except:
            pass

def hdf5_model_to_binary_ses3d_model(input_filename, output_folder):
    with h5py.File(input_filename, "r") as f:
        _hdf5_model_to_binary_ses3d_model(f=f,
                                          output_folder=output_folder)

def _hdf5_model_to_binary_ses3d_model(f, output_folder):
    lpd = 4

    assert not os.path.exists(output_folder), \
        "Folder '%s' already exists." % output_folder
    os.makedirs(output_folder)

    with io.BytesIO(f["_meta"]["boxfile"].value.tostring()) as buf:
        setup = _read_boxfile(buf)
        # Also write to output folder
        buf.seek(0, 0)
        with io.open(os.path.join(output_folder, "boxfile"), "wb") as fh:
            fh.write(buf.read())

    data = {}

    # SES3D internally expects a density in kg/m^3 - The hdf5 files might
    # have g/cm^3.
    rho = f["data"]["rho"][:]
    if rho.mean() < 1000:
        rho *= 1000.0
    data["rhoinv"] = 1.0 / rho

    data["mu"] = (f["data"]["vsh"][:] * 1000) ** 2 / data["rhoinv"]
    data["lambda"] = \
        (f["data"]["vp"][:] * 1000) ** 2 / data["rhoinv"] - 2 * data["mu"]
    data["A"] = (f["data"]["A"][:])
    data["B"] = (f["data"]["vsv"][:] * 1000) ** 2 / data["rhoinv"] - data["mu"]
    data["C"] = (f["data"]["C"][:])
    # Q might now always be given.
    if "Q" in f["data"]:
        data["Q"] = (f["data"]["Q"][:])

    for key in sorted(data.keys()):
        for number, domain in enumerate(setup["subdomains"]):
            x_min, x_max = domain["boundaries_x"]
            y_min, y_max = domain["boundaries_y"]
            z_min, z_max = domain["boundaries_z"]

            # Minimum indices
            x_min, y_min, z_min = \
                [lpd * _j for _j in (x_min, y_min, z_min)]
            # Maximum indices
            x_max, y_max, z_max = \
                [lpd * (_j + 1) for _j in (x_max, y_max, z_max)]

            _d = data[key][x_min: x_max + 1,
                           y_min: y_max + 1,
                           z_min: z_max + 1]
            # Invert last components.
            _d = _d[:, :, ::-1]
            # Reduplicate the GLL points.
            for _i in xrange(3):
                _s = _d.shape[_i]
                left_idx = np.arange(_s - lpd)[::lpd]
                right_idx = np.arange(_s + lpd)[lpd + 1::lpd]
                if _i == 0:
                    _t = [_d[_l:_r, :, :]
                          for _l, _r in zip(left_idx, right_idx)]
                elif _i == 1:
                    _t = [_d[:, _l:_r, :]
                          for _l, _r in zip(left_idx, right_idx)]
                elif _i == 2:
                    _t = [_d[:, :, _l:_r]
                          for _l, _r in zip(left_idx, right_idx)]
                else:
                    raise NotImplementedError

                _d = np.concatenate(_t, axis=_i)
            # Reshape to restore 6 dimensional layout.
            _d = np.require(_d, requirements=["C_CONTIGUOUS"])
            shape = (domain["index_x_count"], lpd + 1,
                     domain["index_y_count"], lpd + 1,
                     domain["index_z_count"], lpd + 1)
            _d = _d.reshape(shape, order="C")
            # Roll to retrieve original SES3D memory order.
            _d = np.rollaxis(_d, 2, 1)
            _d = np.rollaxis(_d, 4, 2)
            _d = np.require(_d, requirements=["F_CONTIGUOUS"])

            filename = os.path.join(output_folder, "%s%i" % (key, number))
            with io.open(filename, "wb") as fh:
                fh.write(struct.pack("<I", 520000))
                fh.write(_d.tobytes(order="F"))
                fh.write(struct.pack("<I", 520000))

def _read_boxfile(fh):
    """
    Copied straight from LASIF.
    """
    setup = {"subdomains": []}

    # The first 14 lines denote the header
    lines = fh.readlines()[14:]
    # Strip lines and remove empty lines.
    lines = [_i.strip() for _i in lines if _i.strip()]

    # The next 4 are the global CPU distribution.
    setup["total_cpu_count"] = int(lines.pop(0))
    setup["cpu_count_in_x_direction"] = int(lines.pop(0))
    setup["cpu_count_in_y_direction"] = int(lines.pop(0))
    setup["cpu_count_in_z_direction"] = int(lines.pop(0))

    if set(lines[0]) == set("-"):
        lines.pop(0)
    # Small sanity check.
    if setup["total_cpu_count"] != setup["cpu_count_in_x_direction"] * \
            setup["cpu_count_in_y_direction"] * \
            setup["cpu_count_in_z_direction"]:
        msg = ("Invalid boxfile. Total and individual processor "
               "counts do not match.")
        raise ValueError(msg)

    # Now parse the rest of file which contains the subdomains.
    def subdomain_generator(data):
        """
        Simple generator looping over each defined box and yielding
        a dictionary for each.

        :param data: The text.
        """
        while data:
            subdom = {}
            # Convert both indices to 0-based indices
            subdom["single_index"] = int(data.pop(0)) - 1
            subdom["multi_index"] = map(lambda x: int(x) - 1,
                                        data.pop(0).split())
            subdom["boundaries_x"] = map(int, data.pop(0).split())
            subdom["boundaries_y"] = map(int, data.pop(0).split())
            subdom["boundaries_z"] = map(int, data.pop(0).split())
            # Convert radians to degree.
            subdom["physical_boundaries_x"] = map(
                    lambda x: math.degrees(float(x)), data.pop(0).split())
            subdom["physical_boundaries_y"] = map(
                    lambda x: math.degrees(float(x)), data.pop(0).split())
            # z is in meter.
            subdom["physical_boundaries_z"] = \
                map(float, data.pop(0).split())
            for component in ("x", "y", "z"):
                idx = "boundaries_%s" % component
                index_count = subdom[idx][1] - subdom[idx][0] + 1
                subdom["index_%s_count" % component] = index_count
                # The boxfiles are slightly awkward in that the indices
                # are not really continuous. For example if one box
                # has 22 as the last index, the first index of the next
                # box will also be 22, even though it should be 23. The
                # next snippet attempts to fix this deficiency.
                offset = int(round(subdom[idx][0] /
                                   float(index_count - 1)))
                subdom[idx][0] += offset
                subdom[idx][1] += offset
            # Remove separator_line if existent.
            if set(lines[0]) == set("-"):
                lines.pop(0)
            yield subdom
    # Sort them after with the single index.
    setup["subdomains"] = sorted(list(subdomain_generator(lines)),
                                 key=lambda x: x["single_index"])
    # Do some more sanity checks.
    if len(setup["subdomains"]) != setup["total_cpu_count"]:
        msg = ("Invalid boxfile. Number of processors and subdomains "
               "to not match.")
        raise ValueError(msg)
    for component in ("x", "y", "z"):
        idx = "index_%s_count" % component
        if len(set([_i[idx] for _i in setup["subdomains"]])) != 1:
            msg = ("Invalid boxfile. Unequal %s index count across "
                   "subdomains.") % component
            raise ValueError(msg)

    # Now generate the absolute indices for the whole domains.
    for component in ("x", "y", "z"):
        setup["boundaries_%s" % component] = (
            min([_i["boundaries_%s" % component][0]
                 for _i in setup["subdomains"]]),
            max([_i["boundaries_%s" %
                    component][1] for _i in setup["subdomains"]]))
        setup["physical_boundaries_%s" % component] = (
            min([_i["physical_boundaries_%s" % component][0] for
                 _i in setup["subdomains"]]),
            max([_i["physical_boundaries_%s" % component][1] for _i in
                 setup["subdomains"]]))

    return setup

def plot_hdf5_model(filename, plot_type="horizontal", *args, **kwargs):
    with h5py.File(filename, "r") as f:
        if plot_type == "horizontal":
            _plot_hdf5_model_horizontal(f=f, *args, **kwargs)
        elif plot_type == "vertical":
            _plot_hdf5_model_vertical(f=f, *args, **kwargs)
        else:
            raise NotImplementedError

def _plot_hdf5_model_vertical(f, component, output_filename, vmin=None,
                              vmax=None):
    import matplotlib.cm
    import matplotlib.pylab as plt

    data = xarray.DataArray(
        f["data"][component][:], [
            ("latitude", 90.0 - f["coordinate_0"][:]),
            ("longitude", f["coordinate_1"][:]),
            ("radius", f["coordinate_2"][:] / 1000.0)])

    plt.style.use('seaborn-pastel')

    plt.figure(figsize=(32, 18))

    plt.suptitle("Component %s - File %s" % (component, output_filename),
                 fontsize=20)

    count = 12
    lats = plt.linspace(data["latitude"].min(), data["latitude"].max(),
                        count)
    lngs = plt.linspace(data["longitude"].min(), data["longitude"].max(),
                        count)

    import lasif.colors
    my_colormap = lasif.colors.get_colormap(
        "tomo_full_scale_linear_lightness")

    # Overwrite colormap things if given.
    if vmin is not None and vmax is not None:
        min_val_plot = vmin
        max_val_plot = vmax
    else:
        mean = data.mean()
        max_diff = max(abs(mean - data.min()),
                       abs(data.max() - mean))
        min_val_plot = mean - max_diff
        max_val_plot = mean + max_diff
        # Plotting essentially constant models.
        min_delta = 0.001 * abs(max_val_plot)
        if (max_val_plot - min_val_plot) < min_delta:
            max_val_plot = max_val_plot + min_delta
            min_val_plot = min_val_plot - min_delta

    for _i in range(count):
        plt.subplot(4, count // 2, _i + 1)

        x, y = np.meshgrid(data.longitude, data.radius)

        plot_data = data.sel(latitude=lats[_i], method="nearest")
        plot_data = np.ma.masked_invalid(plot_data.data)

        # Plot.
        plt.pcolormesh(
            x, y, plot_data.T,
            cmap=my_colormap, vmin=min_val_plot, vmax=max_val_plot,
            shading="flat")

        # make a colorbar and title
        plt.colorbar()
        plt.title("@Latitude: " + str(lats[_i]))

    for _i in range(count):
        plt.subplot(4, count // 2, count + _i + 1)

        x, y = np.meshgrid(data.latitude, data.radius)

        plot_data = data.sel(longitude=lngs[_i], method="nearest")
        plot_data = np.ma.masked_invalid(plot_data.data)

        # Plot.
        plt.pcolormesh(
            x, y, plot_data.T,
            cmap=my_colormap, vmin=min_val_plot, vmax=max_val_plot,
            shading="flat")

        # make a colorbar and title
        plt.colorbar()
        plt.title("@Longitude: " + str(lngs[_i]))

    plt.tight_layout(rect=(0, 0, 1, 0.95))
    plt.savefig(output_filename, dpi=150)
    plt.close()

def _plot_hdf5_model_horizontal(f, component, output_filename,
                                vmin=None, vmax=None):
    import matplotlib.cm
    import matplotlib.pylab as plt

    data = xarray.DataArray(
        f["data"][component][:], [
            ("latitude", 90.0 - f["coordinate_0"][:]),
            ("longitude", f["coordinate_1"][:]),
            ("radius", f["coordinate_2"][:] / 1000.0)])

    plt.style.use('seaborn-pastel')

    from lasif.domain import RectangularSphericalSection
    domain = RectangularSphericalSection(**dict(f["_meta"]["domain"].attrs))

    plt.figure(figsize=(32, 18))

    depth_position_map = {
        50: (0, 0),
        100: (0, 1),
        150: (1, 0),
        250: (1, 1),
        400: (2, 0),
        600: (2, 1)
    }

    for depth, location in depth_position_map.items():
        ax = plt.subplot2grid((3, 5), location)
        radius = 6371.0 - depth

        # set up a map and colourmap
        m = domain.plot(ax=ax, resolution="c", skip_map_features=True)

        import lasif.colors
        my_colormap = lasif.colors.get_colormap(
                "tomo_full_scale_linear_lightness")

        from lasif import rotations

        x, y = np.meshgrid(data.longitude, data.latitude)

        x_shape = x.shape
        y_shape = y.shape

        lat_r, lon_r = rotations.rotate_lat_lon(
                y.ravel(), x.ravel(),
                domain.rotation_axis,
                domain.rotation_angle_in_degree)

        x, y = m(lon_r, lat_r)

        x.shape = x_shape
        y.shape = y_shape

        plot_data = data.sel(radius=radius, method="nearest")
        plot_data = np.ma.masked_invalid(plot_data.data)

        # Overwrite colormap things if given.
        if vmin is not None and vmax is not None:
            min_val_plot = vmin
            max_val_plot = vmax
        else:
            mean = plot_data.mean()
            max_diff = max(abs(mean - plot_data.min()),
                           abs(plot_data.max() - mean))
            min_val_plot = mean - max_diff
            max_val_plot = mean + max_diff
            # Plotting essentially constant models.
            min_delta = 0.001 * abs(max_val_plot)
            if (max_val_plot - min_val_plot) < min_delta:
                max_val_plot = max_val_plot + min_delta
                min_val_plot = min_val_plot - min_delta

        # Plot.
        im = m.pcolormesh(
                x, y, plot_data,
                cmap=my_colormap, vmin=min_val_plot, vmax=max_val_plot,
                shading="gouraud")

        # make a colorbar and title
        m.colorbar(im, "right", size="3%", pad='2%')
        plt.title(str(depth) + ' km')

    # Depth based statistics.
    plt.subplot2grid((3, 5), (0, 4), rowspan=3)
    plt.title("Depth statistics")
    mean = data.mean(axis=(0, 1))
    std = data.std(axis=(0, 1))
    _min = data.min(axis=(0, 1))
    _max = data.max(axis=(0, 1))

    plt.fill_betweenx(data.radius, mean - std, mean + std,
                      label="std", color="#FF3C83")
    plt.plot(mean, data.radius, label="mean", color="k", lw=2)
    plt.plot(_min, data.radius, color="grey", label="min")
    plt.plot(_max, data.radius, color="grey", label="max")
    plt.legend(loc="best")
    plt.xlabel("Value")
    plt.ylabel("Radius")

    plt.hlines(data.radius, plt.xlim()[0], plt.xlim()[1], color="0.8",
               zorder=-10, linewidth=0.5)

    # Roughness plots.
    plt.subplot2grid((3, 5), (0, 2))
    _d = np.abs(data.diff("latitude", n=1)).sum("latitude").data
    plt.title("Roughness in latitude direction, Total: %g" % _d.sum())
    plt.pcolormesh(data.longitude.data, data.radius.data,
                   _d.T, cmap=matplotlib.cm.viridis)
    try:
        plt.colorbar()
    except:
        pass
    plt.xlabel("Longitude")
    plt.ylabel("Radius")

    plt.subplot2grid((3, 5), (1, 2))
    _d = np.abs(data.diff("longitude", n=1)).sum("longitude").data
    plt.title("Roughness in longitude direction. Total: %g" % data.sum())
    plt.pcolormesh(data.latitude.data, data.radius.data, _d.T,
                   cmap=matplotlib.cm.viridis)
    try:
        plt.colorbar()
    except:
        pass
    plt.xlabel("Latitude")
    plt.ylabel("Radius")

    plt.subplot2grid((3, 5), (2, 2))
    _d = np.abs(data.diff("radius", n=1)).sum("radius").data
    plt.title("Roughness in radius direction. Total: %g" % _d.sum())
    plt.pcolormesh(data.longitude.data, data.latitude.data,
                   _d, cmap=matplotlib.cm.viridis)
    try:
        plt.colorbar()
    except:
        pass
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")

    # L2
    plt.subplot2grid((3, 5), (0, 3))
    _d = (data ** 2).sum("latitude").data
    plt.title("L2 Norm in latitude direction, Total: %g" % _d.sum())
    plt.pcolormesh(data.longitude.data, data.radius.data,
                   _d.T, cmap=matplotlib.cm.viridis)
    try:
        plt.colorbar()
    except:
        pass
    plt.xlabel("Longitude")
    plt.ylabel("Radius")

    plt.subplot2grid((3, 5), (1, 3))
    _d = (data ** 2).sum("longitude").data
    plt.title("L2 Norm in longitude direction, Total: %g" % _d.sum())
    plt.pcolormesh(data.latitude.data, data.radius.data, _d.T,
                   cmap=matplotlib.cm.viridis)
    try:
        plt.colorbar()
    except:
        pass
    plt.xlabel("Latitude")
    plt.ylabel("Radius")

    plt.subplot2grid((3, 5), (2, 3))
    _d = (data ** 2).sum("radius").data
    plt.title("L2 Norm in radius direction, Total: %g" % _d.sum())
    plt.pcolormesh(data.longitude.data, data.latitude.data,
                   _d, cmap=matplotlib.cm.viridis)
    try:
        plt.colorbar()
    except:
        pass
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")

    plt.suptitle("Component %s - File %s" % (component, output_filename),
                 fontsize=20)

    plt.tight_layout(rect=(0, 0, 1, 0.95))

    plt.savefig(output_filename, dpi=150)
    plt.close()

def taper_and_precondition_hdf5_model(
        input_filename, output_filename, taper_colatitude_offset_in_km,
        taper_colatitude_width_in_km, taper_longitude_offset_in_km,
        taper_longitude_width_in_km, taper_depth_offset_in_km,
        taper_depth_width_in_km, scaling_file):
    # Make a copy of the file and then modify in-place.
    assert not os.path.exists(output_filename), "File '%s' already exists." % \
        output_filename

    shutil.copy2(input_filename, output_filename)

    with h5py.File(output_filename, "r+") as f:
        _taper_and_precondition_hdf5_model(
            f=f,
            taper_colatitude_offset_in_km=taper_colatitude_offset_in_km,
            taper_colatitude_width_in_km=taper_colatitude_width_in_km,
            taper_longitude_offset_in_km=taper_longitude_offset_in_km,
            taper_longitude_width_in_km=taper_longitude_width_in_km,
            taper_depth_offset_in_km=taper_depth_offset_in_km,
            taper_depth_width_in_km=taper_depth_width_in_km,
            scaling_file=scaling_file)

def _taper_and_precondition_hdf5_model(f, taper_colatitude_offset_in_km,
                                       taper_colatitude_width_in_km,
                                       taper_longitude_offset_in_km,
                                       taper_longitude_width_in_km,
                                       taper_depth_offset_in_km,
                                       taper_depth_width_in_km,
                                       scaling_file):

    # Read the scaling file and make sure it plays nice with the gradient at
    # hand.
    with io.open(scaling_file, "rb") as fh:
        scaling = json.load(fh)
    np.testing.assert_allclose(scaling["radius"], f["coordinate_2"][:])
    scaling = np.array(scaling["weights"], dtype=np.float32)

    fac = 111.19492664455873
    colatitude_in_km = f["coordinate_0"][:] * fac
    longitude_in_km = f["coordinate_1"][:] * fac
    radius_in_km = f["coordinate_2"][:] / 1000.0

    # Convert into distance from either end.
    for _i in [colatitude_in_km, longitude_in_km]:
        _i[:] = np.fmin(_i - _i.min(), _i.max() - _i)
    # In the radial direction we only taper at the bottom.
    radius_in_km -= radius_in_km.min()

    # Apply the offsets
    colatitude_in_km -= taper_colatitude_offset_in_km
    longitude_in_km -= taper_longitude_offset_in_km
    radius_in_km -= taper_depth_offset_in_km

    # Apply the taper width
    colatitude_in_km /= taper_colatitude_width_in_km
    longitude_in_km /= taper_longitude_width_in_km
    radius_in_km /= taper_depth_width_in_km

    # Clip
    longitude_in_km = longitude_in_km.clip(min=0.0, max=1.0)
    colatitude_in_km = colatitude_in_km.clip(min=0.0, max=1.0)
    radius_in_km = radius_in_km.clip(min=0.0, max=1.0)

    # Apply Hanning taper. This finalizes the taper we have to multiply the
    # data with.
    for x in [longitude_in_km, colatitude_in_km, radius_in_km]:
        x[:] = 0.5 * (1.0 - np.cos(x * np.pi))

    # Apply the tapers.
    for name, data in f["data"].items():
        data = data[:]
        data *= colatitude_in_km[:, np.newaxis, np.newaxis]
        data *= longitude_in_km[np.newaxis, :, np.newaxis]
        data *= radius_in_km[np.newaxis, np.newaxis, :]
        # Apply the depth weighting.
        data *= scaling[np.newaxis, np.newaxis, :]
        f["data"][name][:] = data

def determine_depth_scaling(input_filename, output_filename, max_kernel_value):
    with h5py.File(input_filename, mode="r") as f:
        _determine_depth_scaling(f=f,
                                 output_filename=output_filename,
                                 max_kernel_value=max_kernel_value)

def _determine_depth_scaling(f, output_filename, max_kernel_value):
    all_scales = []

    for data in f["data"].values():
        data = data[:]
        # Zeros mess with everything - replace with the smallest
        # non-zero number!
        data[data == 0] = np.abs(data[data != 0]).min()

        # Damping factor - the higher the damping the lesser the effect of
        # the depth scaling.
        damp = 0.1

        m = np.max(np.abs(data))
        fac = np.zeros(data.shape[-1])
        for _i in range(len(fac)):
            fac[_i] = 1.0 / (damp * m + np.abs(data[:, :, _i]).max())

        all_scales.append(fac)

    s = np.sum(all_scales, axis=0)

    import scipy.signal
    # Smooth a tiny bit to avoid wild oscillations.
    w = scipy.signal.gaussian(5, 3)
    w /= w.sum()

    # Scale this for funsies.
    s /= s.min()

    # Avoid boundary effects.
    l = len(s)
    s = np.concatenate([np.ones_like(s) * s[0], s, np.ones_like(s) * s[-1]])

    smooth_s = np.convolve(s, w, mode="same")
    # Cut out the original segment.
    smooth_s = smooth_s[l:-l]
    s = s[l:-l]

    # Abuse ObsPy to taper a bit at both ends.
    smooth_s = obspy.Trace(data=smooth_s).taper(
        max_percentage=0.2, type="cosine", side="left").data.clip(min=1.0)

    # Get the max absolute value in depth for the vsv kernel.
    max_vsv = np.abs(f["data"]["vsv"][:]).max(axis=(0, 1))

    factor = max_kernel_value / (smooth_s * max_vsv).max()

    import matplotlib.pyplot as plt
    plt.style.use("ggplot")

    y = f["coordinate_2"][:] / 1000.0

    plt.subplot(141)
    m = max_vsv
    plt.plot(m, y)
    plt.xlim(-0.1 * m.ptp(), 1.1 * m.max())
    plt.ylim(y[0], y[-1])
    plt.title("max abs vsv")

    plt.subplot(142)
    plt.plot(s, y)
    plt.ylim(y[0], y[-1])
    plt.title("raw")

    plt.subplot(143)
    plt.plot(smooth_s, y)
    plt.xlim(0, smooth_s.max() * 1.5)
    plt.ylim(y[0], y[-1])
    plt.title("smoothed")

    plt.subplot(144)
    m = smooth_s * factor * max_vsv
    plt.plot(m, y)
    plt.xlim(-0.1 * max_kernel_value, 1.1 * max_kernel_value)
    plt.ylim(y[0], y[-1])
    plt.title("after")

    plt.suptitle("Factor: %s" % str(factor))

    output = {
        "radius": [float(i) for i in f["coordinate_2"][:]],
        "weights": [float(i) for i in smooth_s * factor]
    }

    with io.open(output_filename, "wb") as fh:
        json.dump(output, fh)

    plt.show()