<filename>env.py
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
from scipy.io import loadmat, savemat
import networkx as nx
import numpy as np
# import matplotlib.pyplot as plt

class Topology(object):
    def __init__(self, config, data_dir='./data/'):
        self.topology_file = data_dir + 'topo/' + config.topology_file
        self.shortest_paths_file = self.topology_file + '_shortest_paths'
        self.DG = nx.DiGraph()

        self.load_topology()
        self.calculate_paths()

    def load_topology(self):
        print('[*] Loading topology...', self.topology_file)
        f = open(self.topology_file, 'r')
        header = f.readline()
        header = header.split('\t')
        self.num_nodes = int(header[1])
        self.num_links = int(header[3])
        f.readline()
        self.link_idx_to_sd = {}
        self.link_sd_to_idx = {}
        self.link_capacities = np.empty((self.num_links))
        self.link_weights = np.empty((self.num_links))
        for line in f:
            link = line.split('\t')
            i, s, d, w, c = link
            self.link_idx_to_sd[int(i)] = (int(s), int(d))
            self.link_sd_to_idx[(int(s), int(d))] = int(i)
            self.link_capacities[int(i)] = float(c)
            self.link_weights[int(i)] = int(w)
            self.DG.add_weighted_edges_from([(int(s), int(d), int(w))])

        assert len(self.DG.nodes()) == self.num_nodes and len(self.DG.edges()) == self.num_links

        f.close()
        # print('nodes: %d, links: %d\n'%(self.num_nodes, self.num_links))

        # nx.draw_networkx(self.DG)
        # plt.show()

    def calculate_paths(self):
        self.pair_idx_to_sd = []
        self.pair_sd_to_idx = {}
        # Shortest paths
        self.shortest_paths = []
        if os.path.exists(self.shortest_paths_file):
            print('[*] Loading shortest paths...', self.shortest_paths_file)
            f = open(self.shortest_paths_file, 'r')
            self.num_pairs = 0
            for line in f:
                sd = line[:line.find(':')]
                s = int(sd[:sd.find('-')])
                d = int(sd[sd.find('>') + 1:])
                self.pair_idx_to_sd.append((s, d))
                self.pair_sd_to_idx[(s, d)] = self.num_pairs
                self.num_pairs += 1
                self.shortest_paths.append([])
                paths = line[line.find(':') + 1:].strip()[1:-1]
                while paths != '':
                    idx = paths.find(']')
                    path = paths[1:idx]
                    node_path = np.array(path.split(',')).astype(np.int16)
                    assert node_path.size == np.unique(node_path).size
                    self.shortest_paths[-1].append(node_path)
                    paths = paths[idx + 3:]
        else:
            print('[!] Calculating shortest paths...')
            f = open(self.shortest_paths_file, 'w+')
            self.num_pairs = 0
            for s in range(self.num_nodes):
                for d in range(self.num_nodes):
                    if s != d:
                        self.pair_idx_to_sd.append((s, d))
                        self.pair_sd_to_idx[(s, d)] = self.num_pairs
                        self.num_pairs += 1
                        self.shortest_paths.append(list(nx.all_shortest_paths(self.DG, s, d, weight='weight')))
                        line = str(s) + '->' + str(d) + ': ' + str(self.shortest_paths[-1])
                        f.writelines(line + '\n')

        assert self.num_pairs == self.num_nodes * (self.num_nodes - 1)
        f.close()

        print('pairs: %d, nodes: %d, links: %d\n' \
              % (self.num_pairs, self.num_nodes, self.num_links))

class Traffic(object):
    def __init__(self, config, num_nodes, data_dir='../data/', is_training=False):
        self.num_nodes = num_nodes

        splitted_data_fname = os.path.join(data_dir, 'splitted_data/{}.mat'.format(config.dataset))
        if not os.path.isfile(splitted_data_fname):
            self.split_data_from_mat(data_dir=data_dir, dataset=config.dataset)

        print('Load data from ', splitted_data_fname)
        data = loadmat(splitted_data_fname)
        if is_training:
            traffic_matrices = data['train']
        else:
            traffic_matrices = data['test']

        tms_shape = traffic_matrices.shape
        self.tm_cnt = tms_shape[0]
        self.traffic_matrices = np.reshape(traffic_matrices, newshape=(self.tm_cnt, num_nodes, num_nodes))
        self.traffic_file = splitted_data_fname
        print('Traffic matrices dims: [%d, %d, %d]\n' % (self.traffic_matrices.shape[0],
                                                         self.traffic_matrices.shape[1],
                                                         self.traffic_matrices.shape[2]))

    @staticmethod
    def remove_outliers(data):
        q25, q75 = np.percentile(data, 25, axis=0), np.percentile(data, 75, axis=0)
        iqr = q75 - q25
        cut_off = iqr * 3
        lower, upper = q25 - cut_off, q75 + cut_off
        for i in range(data.shape[1]):
            flow = data[:, i]
            flow[flow > upper[i]] = upper[i]
            # flow[flow < lower[i]] = lower[i]
            data[:, i] = flow

        return data

    def train_test_split(self, X, dataset):
        if 'abilene' in dataset:
            train_size = 4 * 7 * 288  # 4 weeks
            val_size = 288 * 7  # 1 week
            test_size = 288 * 7 * 2  # 2 weeks
        elif 'geant' in dataset:
            train_size = 96 * 7 * 4 * 2  # 2 months
            val_size = 96 * 7 * 2  # 2 weeks
            test_size = 96 * 7 * 4  # 1 month
        elif 'brain' in dataset:
            train_size = 1440 * 3  # 3 days
            val_size = 1440  # 1 day
            test_size = 1440 * 2  # 2 days
        elif 'uninett' in dataset:  # granularity: 5 min
            train_size = 4 * 7 * 288  # 4 weeks
            val_size = 288 * 7  # 1 week
            test_size = 288 * 7 * 2  # 2 weeks
        elif 'renater_tm' in dataset:  # granularity: 5 min
            train_size = 4 * 7 * 288  # 4 weeks
            val_size = 288 * 7  # 1 week
            test_size = 288 * 7 * 2  # 2 weeks
        else:
            raise NotImplementedError

        X_train = X[:train_size]

        X_val = X[train_size:val_size + train_size]

        X_test = X[val_size + train_size: val_size + train_size + test_size]

        if 'abilene' in dataset or 'geant' in dataset or 'brain' in dataset:
            X_train = self.remove_outliers(X_train)
            X_val = self.remove_outliers(X_val)

        print('Raw data:')
        print('X_train: ', X_train.shape)
        print('X_val: ', X_val.shape)
        print('X_test: ', X_test.shape)

        return X_train, X_val, X_test

    def split_data_from_mat(self, data_dir, dataset):
        X = self.load_raw(data_dir=data_dir, dataset=dataset)
        train, val, test = self.train_test_split(X, dataset)
        savepath = os.path.join(data_dir, 'splitted_data/')
        if not os.path.exists(savepath):
            os.makedirs(savepath)

        savepathfile = os.path.join(savepath, '{}.mat'.format(dataset))
        savemat(savepathfile, {'train': train,
                               'val': val,
                               'test': test})

    @staticmethod
    def load_raw(data_dir, dataset):
        # load raw data
        data_path = os.path.join(data_dir, 'data/{}.mat'.format(dataset))
        X = loadmat(data_path)['X']
        X = np.reshape(X, newshape=(X.shape[0], -1))
        return X

    def load_traffic_from_mat(self, config):
        pass

    def load_traffic(self, config):
        assert os.path.exists(self.traffic_file)
        print('[*] Loading traffic matrices...', self.traffic_file)
        f = open(self.traffic_file, 'r')
        traffic_matrices = []
        for line in f:
            volumes = line.strip().split(' ')
            total_volume_cnt = len(volumes)
            assert total_volume_cnt == self.num_nodes * self.num_nodes
            matrix = np.zeros((self.num_nodes, self.num_nodes))
            for v in range(total_volume_cnt):
                i = int(v / self.num_nodes)
                j = v % self.num_nodes
                if i != j:
                    matrix[i][j] = float(volumes[v])
            # print(matrix + '\n')
            traffic_matrices.append(matrix)

        f.close()
        self.traffic_matrices = np.array(traffic_matrices)

        tms_shape = self.traffic_matrices.shape
        self.tm_cnt = tms_shape[0]
        print('Traffic matrices dims: [%d, %d, %d]\n' % (tms_shape[0], tms_shape[1], tms_shape[2]))

class Environment(object):
    def __init__(self, config, is_training=False):
        self.data_dir = config.data_dir
        self.topology = Topology(config, self.data_dir)
        self.traffic = Traffic(config, self.topology.num_nodes, self.data_dir, is_training=is_training)
        self.traffic_matrices = self.traffic.traffic_matrices  # kbps
        # self.traffic_matrices = self.traffic.traffic_matrices * 100 * 8 / 300 / 1000  # kbps
        self.tm_cnt = self.traffic.tm_cnt
        self.traffic_file = self.traffic.traffic_file
        self.num_pairs = self.topology.num_pairs
        self.pair_idx_to_sd = self.topology.pair_idx_to_sd
        self.pair_sd_to_idx = self.topology.pair_sd_to_idx
        self.num_nodes = self.topology.num_nodes
        self.num_links = self.topology.num_links
        self.link_idx_to_sd = self.topology.link_idx_to_sd
        self.link_sd_to_idx = self.topology.link_sd_to_idx
        self.link_capacities = self.topology.link_capacities
        self.link_weights = self.topology.link_weights
        self.shortest_paths_node = self.topology.shortest_paths  # paths consist of nodes
        self.shortest_paths_link = self.convert_to_edge_path(self.shortest_paths_node)  # paths consist of links

    def convert_to_edge_path(self, node_paths):
        edge_paths = []
        num_pairs = len(node_paths)
        for i in range(num_pairs):
            edge_paths.append([])
            num_paths = len(node_paths[i])
            for j in range(num_paths):
                edge_paths[i].append([])
                path_len = len(node_paths[i][j])
                for n in range(path_len - 1):
                    e = self.link_sd_to_idx[(node_paths[i][j][n], node_paths[i][j][n + 1])]
                    assert e >= 0 and e < self.num_links
                    edge_paths[i][j].append(e)
                # print(i, j, edge_paths[i][j])

        return edge_paths


<gh_stars>0
from datetime import datetime
import time

import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import interp1d
from scipy.signal import butter, lfilter

__author__ = "<NAME>"
__version__ = "Version 0.9"

measures = {}
working_data = {}

#Data handling
def get_data(filename, delim = ',', column_name = 'None'):
	file_ext = filename.split('.')[-1]
	if file_ext == 'csv' or file_ext == 'txt':
		if (column_name != 'None'):
			hrdata = np.genfromtxt(filename, delimiter=delim, names=True, dtype=None)
			try:
				hrdata = hrdata[column_name]
			except:
				print('\nError loading column "%s" from file "%s". Is column name specified correctly?' %(column_name, filename))
		elif (column_name == 'None'):
			hrdata = np.genfromtxt(filename, delimiter=delim, dtype = np.float64)
		else:
			print('\nError: column name "%s" not found in header of "%s".\n' %(column_name, filename))
	elif file_ext == 'mat':
		print('getting matlab file')
		import scipy.io
		data = scipy.io.loadmat(filename)
		if (column_name != "None"):
			hrdata = np.array(data[column_name][:,0], dtype=np.float64)
		else:
			print("\nError: column name required for Matlab .mat files\n\n")
	else:
		print('unknown file format')
		hrdata = np.nan
	return hrdata

#Preprocessing
def get_samplerate_mstimer(timerdata):
	fs = ((len(timerdata) / (timerdata[-1]-timerdata[0]))*1000)
	working_data['fs'] = fs
	return fs

def get_samplerate_datetime(datetimedata, timeformat='%H:%M:%S.%f'):
	elapsed = ((datetime.strptime(datetimedata[-1], timeformat) - datetime.strptime(datetimedata[0], timeformat)).total_seconds())
	fs = (len(datetimedata) / elapsed)
	working_data['fs'] = fs
	return fs

def rollwindow(x, window):
	shape = x.shape[:-1] + (x.shape[-1] - window + 1, window)
	strides = x.strides + (x.strides[-1],)
	return np.lib.stride_tricks.as_strided(x, shape=shape, strides=strides)

def rolmean(hrdata, hrw, fs):
	avg_hr = (np.mean(hrdata))
	hrarr = np.array(hrdata)
	rol_mean = np.mean(rollwindow(hrarr, int(hrw*fs)), axis=1)
	ln = np.array([avg_hr for i in range(0,abs(len(hrarr)-len(rol_mean))/2)])
	rol_mean = np.insert(rol_mean, 0, ln)
	rol_mean = np.append(rol_mean, ln)
	rol_mean = rol_mean * 1.1
	return rol_mean

def butter_lowpass(cutoff, fs, order=5):
	nyq = 0.5 * fs
	normal_cutoff = cutoff / nyq
	b, a = butter(order, normal_cutoff, btype='low', analog=False)
	return b, a

def butter_lowpass_filter(hrdata, cutoff, fs, order):
	b, a = butter_lowpass(cutoff, fs, order=order)
	y = lfilter(b, a, hrdata)
	return y

def filtersignal(hrdata, cutoff, fs, order):
	hr = np.power(np.array(hrdata), 3)
	hrfiltered = butter_lowpass_filter(hr, cutoff, fs, order)
	return hrfiltered

#Peak detection
def detect_peaks(hrdata, rol_mean, ma_perc, fs):
	rm = np.array(rol_mean)
	rolmean = rm+((rm/100)*ma_perc)
	peaksx = np.where((hrdata > rolmean))[0]
	peaksy = hrdata[np.where((hrdata > rolmean))[0]]
	peakedges = np.concatenate((np.array([0]), (np.where(np.diff(peaksx) > 1)[0]), np.array([len(peaksx)])))
	peaklist = []
	ybeat = []

	for i in range(0, len(peakedges)-1):
		try:
			y = peaksy[peakedges[i]:peakedges[i+1]].tolist()
			peaklist.append(peaksx[peakedges[i] + y.index(max(y))])
		except:
			pass

	working_data['peaklist'] = peaklist
	working_data['ybeat'] = [hrdata[x] for x in peaklist]
	working_data['rolmean'] = rolmean
	calc_RR(fs)
	if len(working_data['RR_list']):
		working_data['rrsd'] = np.std(working_data['RR_list'])
	else:
		working_data['rrsd'] = np.inf

def fit_peaks(hrdata, rol_mean, fs):
	ma_perc_list = [5, 10, 15, 20, 25, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 150, 200, 300]
	rrsd = []
	valid_ma = []
	for x in ma_perc_list:
		detect_peaks(hrdata, rol_mean, x, fs)
		bpm = ((len(working_data['peaklist'])/(len(working_data['hr'])/fs))*60)
		rrsd.append([working_data['rrsd'], bpm, x])

	for x,y,z in rrsd:
		if ((x > 1) and ((y > 40) and (y < 150))):
			valid_ma.append([x, z])

	working_data['best'] = min(valid_ma, key = lambda t: t[0])[1]
	detect_peaks(hrdata, rol_mean, min(valid_ma, key = lambda t: t[0])[1], fs)

def check_peaks():
	RR_arr = np.array(working_data['RR_list'])
	peaklist = np.array(working_data['peaklist'])
	peaklist2 = working_data['peaklist']
	ybeat = np.array(working_data['ybeat'])
	upper_threshold = np.mean(RR_arr) + 300
	lower_threshold = np.mean(RR_arr) - 300
	working_data['RR_list_cor'] = RR_arr[np.where((RR_arr > lower_threshold) & (RR_arr < upper_threshold))]
	peaklist_cor = peaklist[np.where((RR_arr > lower_threshold) & (RR_arr < upper_threshold))[0]+1]
	working_data['peaklist_cor'] = np.insert(peaklist_cor, 0, peaklist[0])
	working_data['removed_beats'] = peaklist[np.where((RR_arr <= lower_threshold) | (RR_arr >= upper_threshold))[0]+1]
	working_data['removed_beats_y'] = removed_beats_y = ybeat[np.where((RR_arr <= lower_threshold) | (RR_arr >= upper_threshold))[0]+1]

#Calculating all measures
def calc_RR(fs):
	peaklist = np.array(working_data['peaklist'])
	RR_list = (np.diff(peaklist) / fs) * 1000.0
	RR_diff = np.abs(np.diff(RR_list))
	RR_sqdiff = np.power(RR_diff, 2)
	working_data['RR_list'] = RR_list
	working_data['RR_diff'] = RR_diff
	working_data['RR_sqdiff'] = RR_sqdiff

def calc_ts_measures():
	RR_list = working_data['RR_list_cor']
	RR_diff = working_data['RR_diff']
	RR_sqdiff = working_data['RR_sqdiff']
	measures['bpm'] = 60000 / np.mean(RR_list)
	measures['ibi'] = np.mean(RR_list)
	measures['sdnn'] = np.std(RR_list)
	measures['sdsd'] = np.std(RR_diff)
	measures['rmssd'] = np.sqrt(np.mean(RR_sqdiff))
	NN20 = [x for x in RR_diff if (x>20)]
	NN50 = [x for x in RR_diff if (x>50)]
	measures['nn20'] = NN20
	measures['nn50'] = NN50
	measures['pnn20'] = float(len(NN20)) / float(len(RR_diff))
	measures['pnn50'] = float(len(NN50)) / float(len(RR_diff))
	measures['hr_mad'] = np.median(np.abs(RR_list-np.median(RR_list)))

def calc_fd_measures(hrdata, fs):
	peaklist = working_data['peaklist_cor']
	RR_list = working_data['RR_list_cor']
	RR_x = peaklist[1:]
	RR_y = RR_list
	RR_x_new = np.linspace(RR_x[0],RR_x[-1],RR_x[-1])
	f = interp1d(RR_x, RR_y, kind='cubic')
	n = len(hrdata)
	frq = np.fft.fftfreq(len(hrdata), d=((1/fs)))
	frq = frq[range(n/2)]
	Y = np.fft.fft(f(RR_x_new))/n
	Y = Y[range(n/2)]
	measures['lf'] = np.trapz(abs(Y[(frq>=0.04) & (frq<=0.15)]))
	measures['hf'] = np.trapz(abs(Y[(frq>=0.16) & (frq<=0.5)]))
	measures['lf/hf'] = measures['lf'] / measures['hf']

#Plotting it
def plotter(show=True, title='Heart Rate Signal Peak Detection'):
	peaklist = working_data['peaklist']
	ybeat = working_data['ybeat']
	rejectedpeaks = working_data['removed_beats']
	rejectedpeaks_y = working_data['removed_beats_y']
	plt.title(title)
	plt.plot(working_data['hr'], alpha=0.5, color='blue', label='heart rate signal')
	plt.scatter(peaklist, ybeat, color='green', label='BPM:%.2f' %(measures['bpm']))
	plt.scatter(rejectedpeaks, rejectedpeaks_y, color='red', label='rejected peaks')
	plt.legend(loc=4, framealpha=0.6)
	if show == True:
		plt.show()
	else:
		return plt

#Wrapper function
def process(hrdata, fs, hrw = 0.75):
	t1 = time.clock()
	hrdata = filtersignal(hrdata, 4, fs, 5)
	working_data['hr'] = hrdata
	print(hrdata)
	rol_mean = rolmean(hrdata, hrw, fs)
	fit_peaks(hrdata, rol_mean, fs)
	calc_RR(fs)
	check_peaks()
	calc_ts_measures()
	calc_fd_measures(hrdata, fs)
	print('\nFinished in %.4fsec' %(time.clock()-t1))
	return measures


<filename>distance/experimental/np_brats.py
import torch
from torch import nn
import torch.nn.functional as F
import numpy as np
from matplotlib import pyplot as plt
from scipy import ndimage
from scipy.ndimage import morphology, distance_transform_edt
from scipy.ndimage import _ni_support
from scipy.ndimage.morphology import distance_transform_edt, binary_erosion,\
    generate_binary_structure
from scipy.ndimage.measurements import label, find_objects
from scipy.stats import pearsonr

def surface_distances(result, reference, voxelspacing=None, connectivity=1):
    """
    The distances between the surface voxel of binary objects in result and their
    nearest partner surface voxel of a binary object in reference.
    after: https://github.com/doublechenching/brats_segmentation-pytorch/blob/master/utils/metric/binary.py
    """
    result = np.atleast_1d(result.astype(np.bool))
    reference = np.atleast_1d(reference.astype(np.bool))
    if voxelspacing is not None:
        voxelspacing = _ni_support._normalize_sequence(voxelspacing, result.ndim)
        voxelspacing = np.asarray(voxelspacing, dtype=np.float64)
        if not voxelspacing.flags.contiguous:
            voxelspacing = voxelspacing.copy()

    # binary structure
    footprint = generate_binary_structure(result.ndim, connectivity)

    # test for emptiness
    if 0 == np.count_nonzero(result):
        raise RuntimeError('The first supplied array does not contain any binary object.')
    if 0 == np.count_nonzero(reference):
        raise RuntimeError('The second supplied array does not contain any binary object.')

        # extract only 1-pixel border line of objects
    result_border = result ^ binary_erosion(result, structure=footprint, iterations=1)
    reference_border = reference ^ binary_erosion(reference, structure=footprint, iterations=1)

    # compute average surface distance
    # Note: scipys distance transform is calculated only inside the borders of the
    #       foreground objects, therefore the input has to be reversed
    dt = distance_transform_edt(~reference_border, sampling=voxelspacing)
    sds = dt[result_border]

    return sds

class HausdorffBinaryLoss(nn.Module):
    """
    Hausdorff Distance.

    Computes the (symmetric) Hausdorff Distance (HD) between the binary objects in two
    images. It is defined as the maximum surface distance between the objects.

    Parameters
    ----------
    result : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    reference : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    voxelspacing : float or sequence of floats, optional
        The voxelspacing in a distance unit i.e. spacing of elements
        along each dimension. If a sequence, must be of length equal to
        the input rank; if a single number, this is used for all axes. If
        not specified, a grid spacing of unity is implied.
    connectivity : int
        The neighbourhood/connectivity considered when determining the surface
        of the binary objects. This value is passed to
        `scipy.ndimage.morphology.generate_binary_structure` and should usually be :math:`> 1`.
        Note that the connectivity influences the result in the case of the Hausdorff distance.

    Returns
    -------
    hd : float
        The symmetric Hausdorff Distance between the object(s) in ```result``` and the
        object(s) in ```reference```. The distance unit is the same as for the spacing of
        elements along each dimension, which is usually given in mm.

    See also
    --------
    :func:`assd`
    :func:`asd`

    Notes
    -----
    This is a real metric. The binary images can therefore be supplied in any order.
    after: https://github.com/doublechenching/brats_segmentation-pytorch/blob/master/utils/metric/binary.py
    """
    def __init__(self, voxelspacing=None, connectivity=1):
        self.voxelspacing=voxelspacing
        self.connectivity=connectivity

    def forward(self, result, reference):
        hd1 = surface_distances(result, reference, self.voxelspacing, self.connectivity).max()
        hd2 = surface_distances(reference, result, self.voxelspacing, self.connectivity).max()
        hd = max(hd1, hd2)
        return hd

class Hausdorff95BinaryLoss(nn.Module):
    """
    95th percentile of the Hausdorff Distance.
    Computes the 95th percentile of the (symmetric) Hausdorff Distance (HD) between the binary objects in two
    images. Compared to the Hausdorff Distance, this metric is slightly more stable to small outliers and is
    commonly used in Biomedical Segmentation challenges.
    Parameters as per HausdorffBratsBinaryLoss
    after: https://github.com/doublechenching/brats_segmentation-pytorch/blob/master/utils/metric/binary.py
    """

    def __init__(self, voxelspacing=None, connectivity=1):
        self.voxelspacing=voxelspacing
        self.connectivity=connectivity

    def forward(self, result, reference):
        hd1 = surface_distances(result, reference, self.voxelspacing, self.connectivity)
        hd2 = surface_distances(reference, result, self.voxelspacing, self.connectivity)
        hd95 = np.percentile(np.hstack((hd1, hd2)), 95)
        return hd95

def asd(result, reference, voxelspacing=None, connectivity=1):
    """
    Average surface distance metric.

    Computes the average surface distance (ASD) between the binary objects in two images.

    Parameters
    ----------
    result : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    reference : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    voxelspacing : float or sequence of floats, optional
        The voxelspacing in a distance unit i.e. spacing of elements
        along each dimension. If a sequence, must be of length equal to
        the input rank; if a single number, this is used for all axes. If
        not specified, a grid spacing of unity is implied.
    connectivity : int
        The neighbourhood/connectivity considered when determining the surface
        of the binary objects. This value is passed to
        `scipy.ndimage.morphology.generate_binary_structure` and should usually be :math:`> 1`.
        The decision on the connectivity is important, as it can influence the results
        strongly. If in doubt, leave it as it is.

    Returns
    -------
    asd : float
        The average surface distance between the object(s) in ``result`` and the
        object(s) in ``reference``. The distance unit is the same as for the spacing
        of elements along each dimension, which is usually given in mm.

    See also
    --------
    :func:`assd`
    :func:`hd`

    Notes
    -----
    This is not a real metric, as it is directed. See `assd` for a real metric of this.

    The method is implemented making use of distance images and simple binary morphology
    to achieve high computational speed.

    after: https://github.com/doublechenching/brats_segmentation-pytorch/blob/master/utils/metric/binary.py
    """
    sds = surface_distances(result, reference, voxelspacing, connectivity)
    asd = sds.mean()
    return asd

def assd(result, reference, voxelspacing=None, connectivity=1):
    """
    Average symmetric surface distance.

    Computes the average symmetric surface distance (ASD) between the binary objects in
    two images.

    Parameters
    ----------
    result : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    reference : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    voxelspacing : float or sequence of floats, optional
        The voxelspacing in a distance unit i.e. spacing of elements
        along each dimension. If a sequence, must be of length equal to
        the input rank; if a single number, this is used for all axes. If
        not specified, a grid spacing of unity is implied.
    connectivity : int
        The neighbourhood/connectivity considered when determining the surface
        of the binary objects. This value is passed to
        `scipy.ndimage.morphology.generate_binary_structure` and should usually be :math:`> 1`.
        The decision on the connectivity is important, as it can influence the results
        strongly. If in doubt, leave it as it is.

    Returns
    -------
    assd : float
        The average symmetric surface distance between the object(s) in ``result`` and the
        object(s) in ``reference``. The distance unit is the same as for the spacing of
        elements along each dimension, which is usually given in mm.

    See also
    --------
    :func:`asd`
    :func:`hd`

    Notes
    -----
    This is a real metric, obtained by calling and averaging
    after: https://github.com/doublechenching/brats_segmentation-pytorch/blob/master/utils/metric/binary.py
    """
    assd = np.mean(
        (asd(result, reference, voxelspacing, connectivity), asd(reference, result, voxelspacing, connectivity)))
    return assd

class AverageSymmetricSurfaceDistance():

    def __init__(self, voxelspacing=None, connectivity=1):
        super(AverageSymmetricSurfaceDistance, self).__init__()
        self.voxelspacing=voxelspacing
        self.connectivity=connectivity

    def forward(self, net_output, target):
        a = assd(net_output.detach().cpu().numpy(), target.detach().cpu().numpy())
        return a.to(net_output.device)

def obj_asd(result, reference, voxelspacing=None, connectivity=1):
    """
    Average surface distance between objects.

    First correspondences between distinct binary objects in reference and result are
    established. Then the average surface distance is only computed between corresponding
    objects. Correspondence is defined as unique and at least one voxel overlap.

    Parameters
    ----------
    result : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    reference : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    voxelspacing : float or sequence of floats, optional
        The voxelspacing in a distance unit i.e. spacing of elements
        along each dimension. If a sequence, must be of length equal to
        the input rank; if a single number, this is used for all axes. If
        not specified, a grid spacing of unity is implied.
    connectivity : int
        The neighbourhood/connectivity considered when determining what accounts
        for a distinct binary object as well as when determining the surface
        of the binary objects. This value is passed to
        `scipy.ndimage.morphology.generate_binary_structure` and should usually be :math:`> 1`.
        The decision on the connectivity is important, as it can influence the results
        strongly. If in doubt, leave it as it is.

    Returns
    -------
    asd : float
        The average surface distance between all mutually existing distinct binary
        object(s) in ``result`` and ``reference``. The distance unit is the same as for the
        spacing of elements along each dimension, which is usually given in mm.

    See also
    --------
    :func:`obj_assd`
    :func:`obj_tpr`
    :func:`obj_fpr`

    Notes
    -----
    This is not a real metric, as it is directed. See `obj_assd` for a real metric of this.
    after: https://github.com/doublechenching/brats_segmentation-pytorch/blob/master/utils/metric/binary.py
    """
    sds = list()
    labelmap1, labelmap2, _a, _b, mapping = __distinct_binary_object_correspondences(result, reference, connectivity)
    slicers1 = find_objects(labelmap1)
    slicers2 = find_objects(labelmap2)
    for lid2, lid1 in list(mapping.items()):
        window = __combine_windows(slicers1[lid1 - 1], slicers2[lid2 - 1])
        object1 = labelmap1[window] == lid1
        object2 = labelmap2[window] == lid2
        sds.extend(surface_distances(object1, object2, voxelspacing, connectivity))
    asd = np.mean(sds)

    return asd

def obj_assd(result, reference, voxelspacing=None, connectivity=1):
    """
    Average symmetric surface distance.

    Computes the average symmetric surface distance (ASSD) between the binary objects in
    two images.

    Parameters
    ----------
    result : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    reference : array_like
        Input data containing objects. Can be any type but will be converted
        into binary: background where 0, object everywhere else.
    voxelspacing : float or sequence of floats, optional
        The voxelspacing in a distance unit i.e. spacing of elements
        along each dimension. If a sequence, must be of length equal to
        the input rank; if a single number, this is used for all axes. If
        not specified, a grid spacing of unity is implied.
    connectivity : int
        The neighbourhood/connectivity considered when determining what accounts
        for a distinct binary object as well as when determining the surface
        of the binary objects. This value is passed to
        `scipy.ndimage.morphology.generate_binary_structure` and should usually be :math:`> 1`.
        The decision on the connectivity is important, as it can influence the results
        strongly. If in doubt, leave it as it is.

    Returns
    -------
    assd : float
        The average symmetric surface distance between all mutually existing distinct
        binary object(s) in ``result`` and ``reference``. The distance unit is the same as for
        the spacing of elements along each dimension, which is usually given in mm.

    after: https://github.com/doublechenching/brats_segmentation-pytorch/blob/master/utils/metric/binary.py
    """
    assd = np.mean((obj_asd(result, reference, voxelspacing, connectivity),
                       obj_asd(reference, result, voxelspacing, connectivity)))
    return assd

class AverageObjectSymmetricSurfaceDistance():
    def __init__(self, voxelspacing=None, connectivity=1):
        super(AverageObjectSymmetricSurfaceDistance, self).__init__()
        self.voxelspacing=voxelspacing
        self.connectivity=connectivity

    def forward(self, net_output, target):
        assd=obj_assd(net_output.detach().cpu().numpy(), target.detach().cpu().numpy())
        return assd.to(net_output.device)

def compute_edts_forPenalizedLoss(GT):
    """
    GT.shape = (batch_size, x,y,z)
    only for binary segmentation
    """
    res = np.zeros(GT.shape)
    for i in range(GT.shape[0]):
        posmask = GT[i]
        negmask = ~posmask
        pos_edt = distance_transform_edt(posmask)
        pos_edt = (np.max(pos_edt) - pos_edt) * posmask
        neg_edt = distance_transform_edt(negmask)
        neg_edt = (np.max(neg_edt) - neg_edt) * negmask

        res[i] = pos_edt / np.max(pos_edt) + neg_edt / np.max(neg_edt)
    return res

def __combine_windows(w1, w2):
    """
    Joins two windows (defined by tuple of slices) such that their maximum
    combined extend is covered by the new returned window.
    after: https://github.com/doublechenching/brats_segmentation-pytorch/blob/master/utils/metric/binary.py
    """
    res = []
    for s1, s2 in zip(w1, w2):
        res.append(slice(min(s1.start, s2.start), max(s1.stop, s2.stop)))

    return tuple(res)

def __distinct_binary_object_correspondences(reference, result, connectivity=1):
    """
    Determines all distinct (where connectivity is defined by the connectivity parameter
    passed to scipy's `generate_binary_structure`) binary objects in both of the input
    parameters and returns a 1to1 mapping from the labelled objects in reference to the
    corresponding (whereas a one-voxel overlap suffices for correspondence) objects in
    result.

    All stems from the problem, that the relationship is non-surjective many-to-many.

    @return (labelmap1, labelmap2, n_lables1, n_labels2, labelmapping2to1)
    after: https://github.com/doublechenching/brats_segmentation-pytorch/blob/master/utils/metric/binary.py
    """
    result = np.atleast_1d(result.astype(np.bool))
    reference = np.atleast_1d(reference.astype(np.bool))

    # binary structure
    footprint = generate_binary_structure(result.ndim, connectivity)

    # label distinct binary objects
    labelmap1, n_obj_result = label(result, footprint)
    labelmap2, n_obj_reference = label(reference, footprint)

    # find all overlaps from labelmap2 to labelmap1; collect one-to-one relationships and store all one-two-many for later processing
    slicers = find_objects(labelmap2)  # get windows of labelled objects
    mapping = dict()  # mappings from labels in labelmap2 to corresponding object labels in labelmap1
    used_labels = set()  # set to collect all already used labels from labelmap2
    one_to_many = list()  # list to collect all one-to-many mappings
    for l1id, slicer in enumerate(slicers):  # iterate over object in labelmap2 and their windows
        l1id += 1  # labelled objects have ids sarting from 1
        bobj = (l1id) == labelmap2[slicer]  # find binary object corresponding to the label1 id in the segmentation
        l2ids = np.unique(labelmap1[slicer][
                                 bobj])  # extract all unique object identifiers at the corresponding positions in the reference (i.e. the mapping)
        l2ids = l2ids[0 != l2ids]  # remove background identifiers (=0)
        if 1 == len(
                l2ids):  # one-to-one mapping: if target label not already used, add to final list of object-to-object mappings and mark target label as used
            l2id = l2ids[0]
            if not l2id in used_labels:
                mapping[l1id] = l2id
                used_labels.add(l2id)
        elif 1 < len(l2ids):  # one-to-many mapping: store relationship for later processing
            one_to_many.append((l1id, set(l2ids)))

    # process one-to-many mappings, always choosing the one with the least labelmap2 correspondences first
    while True:
        one_to_many = [(l1id, l2ids - used_labels) for l1id, l2ids in
                       one_to_many]  # remove already used ids from all sets
        one_to_many = [x for x in one_to_many if x[1]]  # remove empty sets
        one_to_many = sorted(one_to_many, key=lambda x: len(x[1]))  # sort by set length
        if 0 == len(one_to_many):
            break
        l2id = one_to_many[0][1].pop()  # select an arbitrary target label id from the shortest set
        mapping[one_to_many[0][0]] = l2id  # add to one-to-one mappings
        used_labels.add(l2id)  # mark target label as used
        one_to_many = one_to_many[1:]  # delete the processed set from all sets

    return labelmap1, labelmap2, n_obj_result, n_obj_reference, mapping


#!/usr/local/anaconda3/envs/py36 python
# -*- coding: utf-8 -*-

# Plotting
import matplotlib; matplotlib.use('TkAgg')
import matplotlib.pyplot as pl
import seaborn as sns; sns.set_style('ticks')

import matplotlib as mpl
# from matplotlib.ticker import FormatStrFormatter
params = {
    'axes.labelsize': 12,
    'font.size': 12,
    'legend.fontsize': 12,
    'xtick.labelsize': 12,
    'ytick.labelsize': 12,
    'text.usetex': True
}
mpl.rcParams.update(params)

# Imports
from astropy.io import fits
from scipy import interpolate
import numpy as np
import pandas as pd
from util import *
import pyphot
lib = pyphot.get_library()
import astropy.units as u
from scipy.signal import medfilt

def synpassflux(wl, flux, band):
    # Calculate synthetic magnitudes
    n_bands = len(band)

    filt = np.genfromtxt("/Users/jonatanselsing/github/iPTF16geu/data/passbands/%s"%band)
    lamb_T, T = filt[:,0], filt[:,1]
    f = pyphot.Filter(lamb_T, T, name=band, dtype='photon', unit='Angstrom')
    fluxes = f.get_flux(wl, flux, axis=0)
    synmags = -2.5 * np.log10(fluxes) - f.AB_zero_mag
    cwav = np.mean(lamb_T)
    cwave = (float(max(lamb_T[T > np.percentile(T, 10)] - cwav)), float(cwav - min(lamb_T[T > np.percentile(T, 10)])))
    synmag_flux = ((synmags*u.ABmag).to((u.erg/(u.s * u.cm**2 * u.AA)), u.spectral_density(cwav * u.AA))).value
    return synmag_flux, cwav, cwave, synmags

def main():

    root_dir = "../data/"
    OBs = ["OB1", "OB3", "OB4", "OB5", "OB6", "OB7", "OB8"]

    wl_out, flux_out, error_out = [0]*len(OBs), [0]*len(OBs), [0]*len(OBs)
    offsets = [0.95, 0.95, 0.85, 0.95, 0.95, 0.95, 0.95]
    bin_f = 1
    for ii, kk in enumerate(OBs):

        off = 0
        mult = 1.0

        ############################## OB ##############################
        f = fits.open(root_dir + "UVB%s.fits"%kk)
        wl = 10. * f[1].data.field("WAVE").flatten()

        q = f[1].data.field("QUAL").flatten()
        try:
            t = f[1].data.field("TRANS").flatten()
        except:
            t = np.ones_like(q)
        mask_wl = (wl > 3200) & (wl < 5550)
        mask_qual = ~q.astype("bool")
        flux = interpolate.interp1d(wl[mask_qual], f[1].data.field("FLUX").flatten()[mask_qual], bounds_error=False, fill_value=np.nan)

        error = interpolate.interp1d(wl[mask_qual], f[1].data.field("ERR").flatten()[mask_qual], bounds_error=False, fill_value=np.nan)
        wl_plot = wl[mask_wl]
        dust_ext = correct_for_dust(wl_plot, 0.52)
        flux = dust_ext*flux(wl_plot)
        error = dust_ext*error(wl_plot)
        b_wl, b_f, b_e, b_q = bin_spectrum(wl_plot, flux, error, np.zeros_like(flux).astype("bool"), bin_f)

        wl_stitch, flux_stitch, error_stitch = b_wl, b_f, b_e

        f = fits.open(root_dir + "VIS%s.fits"%kk)

        wl = 10. * f[1].data.field("WAVE").flatten()
        q = f[1].data.field("QUAL").flatten()

        try:
            t = f[1].data.field("TRANS").flatten()
        except:
            t = np.ones_like(wl)

        mask_wl = (wl > 5650) & (wl < 10000)
        mask_qual = ~q.astype("bool")

        flux = interpolate.interp1d(wl[mask_qual], f[1].data.field("FLUX").flatten()[mask_qual], bounds_error=False, fill_value=np.nan)
        error = interpolate.interp1d(wl[mask_qual], f[1].data.field("ERR").flatten()[mask_qual], bounds_error=False, fill_value=np.nan)
        wl_plot = wl[mask_wl]
        dust_ext = correct_for_dust(wl_plot, 0.52)
        flux = offsets[ii]*dust_ext*flux(wl_plot) / t[mask_wl]
        error = offsets[ii]*dust_ext*error(wl_plot) / t[mask_wl]
        b_wl, b_f, b_e, b_q = bin_spectrum(wl_plot, flux, error, np.zeros_like(flux).astype("bool"), bin_f)

        wl_stitch, flux_stitch, error_stitch = np.concatenate([wl_stitch,b_wl]), np.concatenate([flux_stitch,b_f]), np.concatenate([error_stitch, b_e])
        # wl_stitch, flux_stitch, error_stitch = np.concatenate([wl_stitch,wl_plot]), np.concatenate([flux_stitch,flux]), np.concatenate([error_stitch, error])

        f = fits.open(root_dir + "NIR%s.fits"%kk)
        wl = 10. * f[1].data.field("WAVE").flatten()
        q = f[1].data.field("QUAL").flatten()

        try:
            t = f[1].data.field("TRANS").flatten()
        except:
            t = np.ones_like(wl)

        mask_wl = (wl > 10000) & (wl < 25000)
        mask_qual = ~q.astype("bool")
        flux = interpolate.interp1d(wl[mask_qual], f[1].data.field("FLUX").flatten()[mask_qual], bounds_error=False, fill_value=np.nan)
        error = interpolate.interp1d(wl[mask_qual], f[1].data.field("ERR").flatten()[mask_qual], bounds_error=False, fill_value=np.nan)
        wl_plot = wl[mask_wl]
        dust_ext = correct_for_dust(wl_plot, 0.52)
        flux = offsets[ii]*dust_ext*flux(wl_plot) / t[mask_wl]
        error = offsets[ii]*dust_ext*error(wl_plot) / t[mask_wl]

        b_wl, b_f, b_e, b_q = bin_spectrum(wl_plot, flux, error, np.zeros_like(flux).astype("bool"), bin_f*3)
        wl_stitch, flux_stitch, error_stitch = np.concatenate([wl_stitch,b_wl]), np.concatenate([flux_stitch,b_f]), np.concatenate([error_stitch, b_e])

        # pl.scatter(wl_stitch, flux_stitch)
        # pl.show()
        wl = np.arange(min(wl_stitch), max(wl_stitch), np.median(np.diff(wl_stitch)))
        f = interpolate.interp1d(wl_stitch, flux_stitch, fill_value=np.nan)
        g = interpolate.interp1d(wl_stitch, error_stitch, fill_value=np.nan)
        # h = interpolate.interp1d(wl_stitch, bp_stitch, fill_value=np.nan)

        wl_out[ii] = wl
        flux_out[ii] = f(wl)
        # flux_out[ii][h(wl) != 0] = np.nan
        error_out[ii] = g(wl)
        # error_out[ii][h(wl) != 0] = np.nan
#
        # np.savetxt(root_dir+"%s_stitched.dat"%kk, list(zip(wl, flux_out[ii], error_out[ii])), fmt='%1.2f %1.4e %1.4e')

    min_wl = 100000
    max_wl = 0
    flux_scaling = [0]*len(OBs)
    for ii, kk in enumerate(OBs):
        min_wl_temp = min(wl_out[ii])
        min_wl = min(min_wl_temp, min_wl)
        max_wl_temp = max(wl_out[ii])
        max_wl = max(max_wl_temp, max_wl)
        dwl = np.median(np.diff(wl_out[ii]))
        flux_scaling[ii] =  np.median(flux_out[ii][(wl_out[ii] > 7000) & (wl_out[ii] < 7100)])
        np.savetxt(root_dir+"%s_stitched.dat"%kk, list(zip(wl, flux_out[ii]*(flux_scaling[0]/flux_scaling[ii]), error_out[ii]*(flux_scaling[0]/flux_scaling[ii]))), fmt='%1.2f %1.4e %1.4e')

    # print(flux_scaling)

    # print(min_wl, max_wl, dwl)
    wl_final = np.arange(min_wl, max_wl, dwl)
    crab_spectra_flux = np.zeros((len(OBs), len(wl_final)))
    crab_spectra_error = np.zeros((len(OBs), len(wl_final)))

    for ii, kk in enumerate(OBs):
        mask = np.isnan(flux_out[ii]).astype("int")
        f = interpolate.interp1d(wl_out[ii], flux_out[ii], bounds_error=False, fill_value=np.nan)
        g = interpolate.interp1d(wl_out[ii], error_out[ii], bounds_error=False, fill_value=np.nan)
        crab_spectra_flux[ii, :] = f(wl_final)*(flux_scaling[0]/flux_scaling[ii])
        crab_spectra_error[ii, :] = g(wl_final)*(flux_scaling[0]/flux_scaling[ii])

        pl.plot(wl_final, crab_spectra_flux[ii, :], color="black", alpha = 0.2, linestyle="steps-mid", rasterized=True)

    print(np.mean(flux_scaling), np.std(flux_scaling))
    print(np.mean(flux_scaling[0]/flux_scaling), np.std(flux_scaling[0]/flux_scaling))

    # l, m, h =(np.percentile(flux_scaling[0]/flux_scaling, [15.9, 50, 84.2]))
    # print(m, m - l, h - m)
    # exit()

    # Median spectrum
    med_spec = np.nanmedian(crab_spectra_flux[:, :], axis=0)
    np.savetxt(root_dir+"median_spectrum.dat", list(zip(wl_final, med_spec)), fmt='%1.2f %1.4e')

    # Weighted mean
    weight = 1.0 / (crab_spectra_error ** 2.0)
    average = np.ma.sum(crab_spectra_flux * weight, axis=0) / np.ma.sum(weight, axis=0)
    std = np.sqrt(np.ma.sum(crab_spectra_error**2 * weight**2, axis=0) / np.ma.sum(weight, axis=0)**2)

    # average =  np.nanmean(crab_spectra_flux, axis=0)
    # std = np.sqrt(np.nansum(crab_spectra_error**2, axis=0))
    # print(average)
    # print(std)

    np.savetxt(root_dir+"weighted_spectrum.dat", list(zip(wl_final, average, std)), fmt='%1.2f %1.4e %1.4e')
    # exit()
    # pl.plot(wl_final, med_spec)
    pl.plot(wl_final, medfilt(average, 1), linestyle="steps-mid", rasterized=True)

    # NUV = np.genfromtxt("../data/crabNUV.txt")
    # dust_ext = correct_for_dust(NUV[:, 0], 0.52)
    # pl.plot(NUV[:, 0], dust_ext*medfilt(NUV[:, 1], 11), linestyle="steps-mid")
    # pl.show()

    # for ii, OB in enumerate(OBs):

    # passbands = [
        # "FORS2_U.dat", "FORS2_B.dat", "FORS2_V.dat", "FORS2_R.dat", "FORS2_I.dat", "2MASS_J.dat", "2MASS_H.dat", "2MASS_Ks.dat"
        # ]
    passbands = ["FORS1_U.dat", "FORS1_B.dat", "FORS1_V.dat", "FORS1_R.dat", "FORS1_I.dat", "FORS2_z.dat", "NACO_J.dat", "NACO_H.dat", "NACO_Ks.dat"]
    photometry = pd.read_csv("../data/Crab_phot.csv")

    vega_conv = pd.read_csv("~/github/iPTF16geu/data/passbands/VEGA_AB.csv")
    # cmap1 = sns.color_palette("viridis", len(passbands))
    cmap1 = sns.color_palette("plasma", 5+len(passbands))

    for pp, ss in enumerate(passbands):
        pass_name = ss.split("_")[-1].split(".")[0]
        meas_mag = photometry.loc[photometry['Bandpass'] == pass_name]
        print(pass_name)
        vega = vega_conv.loc[vega_conv["Band"] == pass_name]
        print(meas_mag["CrabPulsar+Knot"], vega["dmag"].values)
        meas_mag["CrabPulsar+Knot"] = meas_mag["CrabPulsar+Knot"] + vega["dmag"].values
        dust_ext = correct_for_dust(vega["leff"].values, 0.52)
        # print(dust_ext, ebv)

        nanmask = (np.isnan(average) | np.isnan(std)) | (np.isinf(average) | np.isinf(std))
        synmag_flux, cwav, cwave, synmag = synpassflux(wl_final[~nanmask], average[~nanmask], ss)
        synmag_error, cwav, cwave, synmag_err = synpassflux(wl_final[~nanmask], std[~nanmask], ss)
        _, _, _, synmag_err_up = synpassflux(wl_final[~nanmask], average[~nanmask] + std[~nanmask], ss)
        _, _, _, synmag_err_low = synpassflux(wl_final[~nanmask], average[~nanmask] - std[~nanmask], ss)
        e_u = (synmag - synmag_err_up)
        e_l = (synmag_err_low - synmag)
        # pl.errorbar(cwav, synmag_flux, xerr = [[cwave]], yerr = synmag_error,  fmt = 'o', color=cmap1[pp], zorder = 10, ms = 5, elinewidth=1.7, label = "%s = %s$^{+%s}_{-%s}$"%(pass_name, np.around(synmag, 2), np.around(e_u, 2), np.around(e_l, 2)))
        # pl.errorbar(cwav, synmag_flux, xerr = [[cwave]], yerr = synmag_error,  fmt = 'o', color=cmap1[pp], zorder = 10, ms = 5, elinewidth=1.7, label = "%s"%(pass_name))

        if meas_mag.shape[0] > 0:
            meas_flux = ((meas_mag["CrabPulsar+Knot"].values*u.ABmag).to((u.erg/(u.s * u.cm**2 * u.AA)), u.spectral_density(cwav * u.AA))).value
            meas_flux_up =  (((meas_mag["CrabPulsar+Knot"].values + meas_mag["CrabPulsar+Knot_e"].values)*u.ABmag).to((u.erg/(u.s * u.cm**2 * u.AA)), u.spectral_density(cwav * u.AA))).value
            meas_flux_do =  (((meas_mag["CrabPulsar+Knot"].values - meas_mag["CrabPulsar+Knot_e"].values)*u.ABmag).to((u.erg/(u.s * u.cm**2 * u.AA)), u.spectral_density(cwav * u.AA))).value

            pl.errorbar(cwav, dust_ext*meas_flux, xerr = [[cwave]], yerr = [[(meas_flux_do, meas_flux_up)]],  fmt = 'o', color=cmap1[pp], ms = 7, zorder=9, label = "%s"%(pass_name))
    pl.xlim(2900, 25000)
    pl.ylim(1e-16, 2e-14)

    pl.ylabel(r'$\log (F_\lambda / \mathrm{erg}~\mathrm{s}^{-1}~\mathrm{cm}^{-2}~\mathrm{\AA}^{-1}$)')
    pl.semilogy()
    pl.legend()

    # Add frequency axis
    ax = pl.gca()
    ax2 = pl.twiny()

    # get axis limits
    xmin, xmax = ax.get_xlim()
   ax2.set_xlim((xmin, xmax))

    def co(angs):
        return(3e18/(10**angs))
    nu_arr = np.array([15, 14.6, 14.4, 14.3, 14.2, 14.1])
    ax2.set_xticks(co(nu_arr))
    ax2.set_xticklabels(nu_arr)

    ax.set_xlabel(r"$ \lambda_{\mathrm{obs}}/\mathrm{\AA}$")
    ax2.set_xlabel(r"$\log (\nu/\mathrm{Hz})$")

    pl.tight_layout()
    pl.savefig("../figures/combined_spectrum.pdf")
    pl.show()

if __name__ == '__main__':
    main()


r"""Classes and functions to define and manipulate symbolic Operator
expressions. For more details see :ref:`operator_algebra`.

For a list of all properties and methods of an operator object, see the
documentation for the basic :class:`Operator` class.
"""
import re
from abc import ABCMeta, abstractmethod
from collections import OrderedDict, defaultdict
from itertools import product as cartesian_product

from sympy import sympify

from ..pattern_matching import pattern, pattern_head, wc
from ..utils.indices import FockIndex, IdxSym, SymbolicLabelBase
from ..utils.ordering import FullCommutativeHSOrder
from ..utils.properties_for_args import properties_for_args
from ..utils.singleton import Singleton, singleton_object
from .abstract_quantum_algebra import (
    QuantumAdjoint,
    QuantumDerivative,
    QuantumExpression,
    QuantumIndexedSum,
    QuantumOperation,
    QuantumPlus,
    QuantumSymbol,
    QuantumTimes,
    ScalarTimesQuantumExpression,
    SingleQuantumOperation,
    ensure_local_space,
)
from .algebraic_properties import (
    assoc,
    assoc_indexed,
    collect_summands,
    commutator_order,
    delegate_to_method,
    disjunct_hs_zero,
    filter_neutral,
    implied_local_space,
    indexed_sum_over_const,
    indexed_sum_over_kronecker,
    match_replace,
    match_replace_binary,
    orderby,
    scalars_to_op,
)
from .exceptions import CannotSimplify
from .hilbert_space_algebra import (
    HilbertSpace,
    LocalSpace,
    ProductSpace,
    TrivialSpace,
)
from .scalar_algebra import Scalar, ScalarValue, is_scalar

sympyOne = sympify(1)

# for hilbert space dimensions less than or equal to this,
# compute numerically PseudoInverse and NullSpaceProjector representations
DENSE_DIMENSION_LIMIT = 1000

__all__ = [
    'Adjoint',
    'LocalOperator',
    'LocalSigma',
    'NullSpaceProjector',
    'Operator',
    'OperatorPlus',
    'OperatorPlusMinusCC',
    'OperatorSymbol',
    'OperatorTimes',
    'OperatorTrace',
    'PseudoInverse',
    'ScalarTimesOperator',
    'LocalProjector',
    'adjoint',
    'rewrite_with_operator_pm_cc',
    'decompose_space',
    'factor_coeff',
    'factor_for_trace',
    'get_coeffs',
    'II',
    'IdentityOperator',
    'ZeroOperator',
    'OperatorDerivative',
    'Commutator',
    'OperatorIndexedSum',
    'tr',
]

__private__ = []
# anything not in __all__ must be in __private__

###############################################################################
# Abstract base classes
###############################################################################

class Operator(QuantumExpression, metaclass=ABCMeta):
    """Base class for all quantum operators."""

    def pseudo_inverse(self):
        """Pseudo-inverse $\Op{X}^+$ of the operator $\Op{X}$

        It is defined via the relationship

        .. math::

            \Op{X} \Op{X}^+ \Op{X} =  \Op{X}  \\
            \Op{X}^+ \Op{X} \Op{X}^+ =  \Op{X}^+  \\
            (\Op{X}^+ \Op{X})^\dagger = \Op{X}^+ \Op{X}  \\
            (\Op{X} \Op{X}^+)^\dagger = \Op{X} \Op{X}^+
        """
        return self._pseudo_inverse()

    @abstractmethod
    def _pseudo_inverse(self):
        raise NotImplementedError(self.__class__.__name__)

    def expand_in_basis(self, basis_states=None, hermitian=False):
        """Write the operator as an expansion into all
        :class:`KetBras <.KetBra>`
        spanned by `basis_states`.

        Args:
            basis_states (list or None): List of basis states (:class:`.State`
                instances) into which to expand the operator. If None, use the
                operator's `space.basis_states`
            hermitian (bool): If True, assume that the operator is Hermitian
                and represent all elements in the lower triangle of the
                expansion via :class:`OperatorPlusMinusCC`. This is meant to
                enhance readability

        Raises:
            .BasisNotSetError: If `basis_states` is None and the operator's
                Hilbert space has no well-defined basis

        Example:

            >>> hs = LocalSpace(1, basis=('g', 'e'))
            >>> op = LocalSigma('g', 'e', hs=hs) + LocalSigma('e', 'g', hs=hs)
            >>> print(ascii(op, sig_as_ketbra=False))
            sigma_e,g^(1) + sigma_g,e^(1)
            >>> print(ascii(op.expand_in_basis()))
            |e><g|^(1) + |g><e|^(1)
            >>> print(ascii(op.expand_in_basis(hermitian=True)))
            |g><e|^(1) + c.c.
        """
        # fmt: off
        from qalgebra.core.state_algebra import KetBra  # isort:skip
        # KetBra is imported locally to avoid circular imports
        # fmt: on
        if basis_states is None:
            basis_states = list(self.space.basis_states)
        else:
            basis_states = list(basis_states)
        diag_terms = []
        terms = []
        for i, ket_i in enumerate(basis_states):
            for j, ket_j in enumerate(basis_states):
                if i > j and hermitian:
                    continue
                op_ij = (ket_i.dag() * self * ket_j).expand()
                ketbra = KetBra(ket_i, ket_j)
                term = op_ij * ketbra
                if term is not ZeroOperator:
                    if i == j:
                        diag_terms.append(op_ij * ketbra)
                    else:
                        terms.append(op_ij * ketbra)
        if hermitian:
            res = OperatorPlus.create(*diag_terms)
            if len(terms) > 0:
                res = res + OperatorPlusMinusCC(OperatorPlus.create(*terms))
            return res
        else:
            return OperatorPlus.create(*diag_terms) + OperatorPlus.create(
                *terms
            )

class LocalOperator(Operator, metaclass=ABCMeta):
    """Base class for "known" operators on a :class:`.LocalSpace`.

    All :class:`LocalOperator` instances have known algebraic properties and a
    fixed associated identifier (symbol) that is used when printing that
    operator. A custom identifier can be used through the associated
    :class:`.LocalSpace`'s `local_identifiers` parameter. For example::

        >>> hs1_custom = LocalSpace(1, local_identifiers={'Destroy': 'b'})
        >>> b = Destroy(hs=hs1_custom)
        >>> ascii(b)
        'b^(1)'

    Note:
        It is recommended that subclasses use the :func:`.properties_for_args`
        class decorator if they define any position arguments (via the
        ``_arg_names`` class attribute)
    """

    simplifications = [
        implied_local_space(keys=['hs']),
    ]

    _identifier = None  # must be overridden by subclasses!
    _dagger = False  #: do representations include a dagger?
    _arg_names = ()  #: names of args that can be passed to __init__
    _scalar_args = True  #: convert args to Scalar?
    _hs_cls = LocalSpace  #: allowed type of `hs`
    _rx_identifier = re.compile('^[A-Za-z][A-Za-z0-9]*(_[A-Za-z0-9().+-]+)?$')

    def __init__(self, *args, hs):
        if len(args) != len(self._arg_names):
            raise ValueError(
                "expected %d arguments, gotten %d"
                % (len(self._arg_names), len(args))
            )
        if self._scalar_args:
            args = [ScalarValue.create(arg) for arg in args]
        for i, arg_name in enumerate(self._arg_names):
            self.__dict__['_%s' % arg_name] = args[i]
        hs = ensure_local_space(hs, cls=self._hs_cls)
        self._hs = hs
        if self._identifier is None:
            raise TypeError(
                r"Can't instantiate abstract class %s with undefined "
                r"_identifier" % self.__class__.__name__
            )
        self._args = args
        super().__init__(*args, hs=hs)

    @property
    def space(self):
        """Hilbert space of the operator (:class:`.LocalSpace` instance)."""
        return self._hs

    @property
    def args(self):
        """The positional arguments used for instantiating the operator"""
        return tuple(self._args)

    @property
    def kwargs(self):
        """The keyword arguments used for instantiating the operator"""
        return OrderedDict([('hs', self._hs)])

    @property
    def identifier(self):
        """The identifier (symbol) that is used when printing the operator.

        A custom identifier can be used through the associated
        :class:`.LocalSpace`'s `local_identifiers` parameter. For example::

            >>> a = Destroy(hs=1)
            >>> a.identifier
            'a'
            >>> hs1_custom = LocalSpace(1, local_identifiers={'Destroy': 'b'})
            >>> b = Destroy(hs=hs1_custom)
            >>> b.identifier
            'b'
            >>> ascii(b)
            'b^(1)'
        """

        identifier = self._hs._local_identifiers.get(
            self.__class__.__name__, self._identifier
        )
        if not self._rx_identifier.match(identifier):
            raise ValueError(
                "identifier '%s' does not match pattern '%s'"
                % (identifier, self._rx_identifier.pattern)
            )
        return identifier

    def _diff(self, sym):
        return OperatorDerivative(self, derivs={sym: 1})

    def _simplify_scalar(self, func):
        if self._scalar_args:
            args = [arg.simplify_scalar(func=func) for arg in self.args]
            return self.create(*args, hs=self.space)
        else:
            return super()._simplify_scalar(func=func)

###############################################################################
# Operator algebra elements
###############################################################################

class OperatorSymbol(QuantumSymbol, Operator):
    """Symbolic operator

    See :class:`.QuantumSymbol`.
    """

    def _pseudo_inverse(self):
        return PseudoInverse(self)

@singleton_object
class IdentityOperator(Operator, metaclass=Singleton):
    """``IdentityOperator`` constant (singleton) object."""

    _order_index = 2

    @property
    def space(self):
        """:class:`.TrivialSpace`"""
        return TrivialSpace

    @property
    def args(self):
        return tuple()

    def _diff(self, sym):
        return ZeroOperator

    def _adjoint(self):
        return self

    def _pseudo_inverse(self):
        return self

II = IdentityOperator

@singleton_object
class ZeroOperator(Operator, metaclass=Singleton):
    """``ZeroOperator`` constant (singleton) object."""

    _order_index = 2

    @property
    def space(self):
        """:class:`.TrivialSpace`"""
        return TrivialSpace

    @property
    def args(self):
        return tuple()

    def _diff(self, sym):
        return self

    def _adjoint(self):
        return self

    def _pseudo_inverse(self):
        return self

@properties_for_args
class LocalSigma(LocalOperator):
    r"""Level flip operator between two levels of a :class:`.LocalSpace`.

    .. math::

        \Op{\sigma}_{jk}^{\rm hs} =
        \left| j\right\rangle_{\rm hs} \left \langle k \right |_{\rm hs}

    For $j=k$ this becomes a projector $\Op{P}_k$ onto the eigenstate
    $\ket{k}$.

    Args:
        j (int or str): The label or index identifying $\ket{j}$
        k (int or str):  The label or index identifying $\ket{k}$
        hs (LocalSpace or int or str): The Hilbert space on which the
            operator acts. If an :class:`int` or a :class:`str`, an implicit
            Hilbert space will be constructed as a subclass of
            :class:`.LocalSpace`, as configured by :func:`.init_algebra`.

    Note:
        The parameters `j` or `k` may be an integer or a string. A string
        refers to the label of an eigenstate in the basis of `hs`, which needs
        to be set. An integer refers to the (zero-based) index of eigenstate of
        the Hilbert space. This works if `hs` has an unknown dimension.
        Assuming the Hilbert space has a defined basis, using integer or string
        labels is equivalent::

            >>> hs = LocalSpace('tls', basis=('g', 'e'))
            >>> LocalSigma(0, 1, hs=hs) == LocalSigma('g', 'e', hs=hs)
            True

    Raises:
        ValueError: If `j` or `k` are invalid value for the given `hs`

    Printers should represent this operator either in braket notation, or using
    the operator identifier

        >>> LocalSigma(0, 1, hs=0).identifier
        'sigma'

    For ``j == k``, an alternative (fixed) identifier may be used

        >>> LocalSigma(0, 0, hs=0)._identifier_projector
        'Pi'
    """

    _identifier = "sigma"
    _identifier_projector = "Pi"
    _rx_identifier = re.compile('^[A-Za-z][A-Za-z0-9]*$')
    _arg_names = ('j', 'k')
    _scalar_args = False  # args are labels, not scalar coefficients
    _rules = OrderedDict()
    simplifications = [
        implied_local_space(keys=['hs']),
        match_replace,
    ]

    def __init__(self, j, k, *, hs):
        if isinstance(hs, (str, int)):
            hs = self._default_hs_cls(hs)
        hs._unpack_basis_label_or_index(j)  # for applying checks only ...
        hs._unpack_basis_label_or_index(k)  # ... (disregard returned tuple)
        if hs.has_basis:
            # normalize integer i/j to str label, if possible
            if isinstance(j, int):
                j = hs.basis_labels[j]
            if isinstance(k, int):
                k = hs.basis_labels[k]
        super().__init__(j, k, hs=hs)

    @property
    def args(self):
        """The two eigenstate labels `j` and `k` that the operator connects"""
        return self.j, self.k

    @property
    def index_j(self):
        """Index `j` or (zero-based) index of the label `j` in the basis"""
        if isinstance(self.j, (int, SymbolicLabelBase)):
            return self.j
        else:
            try:
                return self.space.basis_labels.index(self.j)
            except ValueError:
                raise ValueError(
                    "%r is not one of the basis labels %r"
                    % (self.j, self.space.basis_labels)
                )

    @property
    def index_k(self):
        """Index `k` or (zero-based) index of the label `k` in the basis"""
        if isinstance(self.k, (int, SymbolicLabelBase)):
            return self.k
        else:
            try:
                return self.space.basis_labels.index(self.k)
            except ValueError:
                raise ValueError(
                    "%r is not one of the basis labels %r"
                    % (self.k, self.space.basis_labels)
                )

    def raise_jk(self, j_incr=0, k_incr=0):
        r"""Return a new :class:`LocalSigma` instance with incremented `j`,
        `k`, on the same Hilbert space:

        .. math::

            \Op{\sigma}_{jk}^{\rm hs} \rightarrow \Op{\sigma}_{j'k'}^{\rm hs}

        This is the result of multiplying $\Op{\sigma}_{jk}^{\rm hs}$
        with any raising or lowering operators.

        If $j'$ or $k'$ are outside the Hilbert space ${\rm hs}$, the result is
        the :obj:`ZeroOperator` .

        Args:
            j_incr (int): The increment between labels $j$ and $j'$
            k_incr (int): The increment between labels $k$ and $k'$. Both
                increments may be negative.
        """
        try:
            if isinstance(self.j, int):
                new_j = self.j + j_incr
            else:  # str
                new_j = self.space.next_basis_label_or_index(self.j, j_incr)
            if isinstance(self.k, int):
                new_k = self.k + k_incr
            else:  # str or SymbolicLabelBase
                new_k = self.space.next_basis_label_or_index(self.k, k_incr)
            return LocalSigma.create(new_j, new_k, hs=self.space)
        except (IndexError, ValueError):
            return ZeroOperator

    def _adjoint(self):
        return LocalSigma(j=self.k, k=self.j, hs=self.space)

    def _pseudo_inverse(self):
        return self._adjoint()

def LocalProjector(j, *, hs):
    """A projector onto a specific level of a :class:`.LocalSpace`

    Args:
        j (int or str): The label or index identifying the state onto which
            is projected
        hs (HilbertSpace): The Hilbert space on which the operator acts
    """
    return LocalSigma(j, j, hs=hs)

###############################################################################
# Algebra Operations
###############################################################################

class OperatorPlus(QuantumPlus, Operator):
    """Sum of Operators"""

    _neutral_element = ZeroOperator
    _binary_rules = OrderedDict()
    simplifications = [
        assoc,
        scalars_to_op,
        orderby,
        collect_summands,
        match_replace_binary,
    ]

    def _pseudo_inverse(self):
        return PseudoInverse(self)

class OperatorTimes(QuantumTimes, Operator):
    """Product of operators

    This serves both as a product within a Hilbert space as well as a tensor
    product."""

    _neutral_element = IdentityOperator
    _binary_rules = OrderedDict()
    simplifications = [assoc, orderby, filter_neutral, match_replace_binary]

    def _pseudo_inverse(self):
        return self.__class__.create(
            *[o._pseudo_inverse() for o in reversed(self.operands)]
        )

class ScalarTimesOperator(Operator, ScalarTimesQuantumExpression):
    """Product of a :class:`.Scalar` coefficient and an :class:`.Operator`."""

    _rules = OrderedDict()
    simplifications = [
        match_replace,
    ]

    @staticmethod
    def has_minus_prefactor(c):
        """
        For a scalar object c, determine whether it is prepended by a "-" sign.
        """
        # TODO: check if this is necessary; if yes, move
        cs = str(c).strip()
        return cs[0] == "-"

    def _pseudo_inverse(self):
        c, t = self.operands
        return t.pseudo_inverse() / c

    def __eq__(self, other):
        # TODO: review, and add this to ScalarTimesQuantumExpression
        if self.term is IdentityOperator and is_scalar(other):
            return self.coeff == other
        return super().__eq__(other)

    def __hash__(self):
        # TODO: review, and add this to ScalarTimesQuantumExpression
        return super().__hash__()

    def _adjoint(self):
        return ScalarTimesOperator(self.coeff.conjugate(), self.term.adjoint())

class OperatorDerivative(QuantumDerivative, Operator):
    """Symbolic partial derivative of an operator

    See :class:`.QuantumDerivative`.
    """

    def _pseudo_inverse(self):
        return PseudoInverse(self)

class Commutator(QuantumOperation, Operator):
    r"""Commutator of two operators

    .. math::

        [\Op{A}, \Op{B}] = \Op{A}\Op{B} - \Op{A}\Op{B}

    """

    _rules = OrderedDict()
    simplifications = [
        scalars_to_op,
        disjunct_hs_zero,
        commutator_order,
        match_replace,
    ]
    # TODO: doit method

    order_key = FullCommutativeHSOrder

    # commutator_order makes FullCommutativeHSOrder anti-commutative

    def __init__(self, A, B):
        self._hs = A.space * B.space
        super().__init__(A, B)

    @property
    def A(self):
        """Left side of the commutator"""
        return self.operands[0]

    @property
    def B(self):
        """Left side of the commutator"""
        return self.operands[1]

    def doit(self, classes=None, recursive=True, **kwargs):
        """Write out commutator

        Write out the commutator according to its definition
        $[\Op{A}, \Op{B}] = \Op{A}\Op{B} - \Op{A}\Op{B}$.

        See :meth:`.Expression.doit`.
        """
        return super().doit(classes, recursive, **kwargs)

    def _doit(self, **kwargs):
        return self.A * self.B - self.B * self.A

    def _expand(self):
        A = self.A.expand()
        B = self.B.expand()
        if isinstance(A, OperatorPlus):
            A_summands = A.operands
        else:
            A_summands = (A,)
        if isinstance(B, OperatorPlus):
            B_summands = B.operands
        else:
            B_summands = (B,)
        summands = []
        for combo in cartesian_product(A_summands, B_summands):
            summands.append(Commutator.create(*combo))
        return OperatorPlus.create(*summands)

    def _series_expand(self, param, about, order):
        A_series = self.A.series_expand(param, about, order)
        B_series = self.B.series_expand(param, about, order)
        res = []
        for n in range(order + 1):
            summands = [
                self.create(A_series[k], B_series[n - k]) for k in range(n + 1)
            ]
            res.append(OperatorPlus.create(*summands))
        return tuple(res)

    def _diff(self, sym):
        return self.__class__(self.A.diff(sym), self.B) + self.__class__(
            self.A, self.B.diff(sym)
        )

    def _adjoint(self):
        return Commutator(self.B.adjoint(), self.A.adjoint())

    def _pseudo_inverse(self):
        return PseudoInverse(self)

class OperatorTrace(SingleQuantumOperation, Operator):
    r"""(Partial) trace of an operator

    Trace of an operator `op` ($\Op{O}) over the degrees
    of freedom of a Hilbert space `over_space` ($\mathcal{H}$):

    .. math::

        {\rm Tr}_{\mathcal{H}} \Op{O}

    Args:
        over_space (.HilbertSpace): The degrees of freedom to trace over
        op (Operator): The operator to take the trace of.
    """
    _rules = OrderedDict()
    simplifications = [
        scalars_to_op,
        implied_local_space(keys=['over_space']),
        match_replace,
    ]

    def __init__(self, op, *, over_space):
        if isinstance(over_space, (int, str)):
            over_space = self._default_hs_cls(over_space)
        assert isinstance(over_space, HilbertSpace)
        self._over_space = over_space
        super().__init__(op, over_space=over_space)
        self._space = None

    @property
    def kwargs(self):
        return {'over_space': self._over_space}

    @property
    def operand(self):
        return self.operands[0]

    @property
    def space(self):
        if self._space is None:
            return self.operands[0].space / self._over_space
        return self._space

    def _expand(self):
        s = self._over_space
        o = self.operand
        return OperatorTrace.create(o.expand(), over_space=s)

    def _series_expand(self, param, about, order):
        ope = self.operand.series_expand(param, about, order)
        return tuple(
            OperatorTrace.create(opet, over_space=self._over_space)
            for opet in ope
        )

    def _diff(self, sym):
        s = self._over_space
        o = self.operand
        return OperatorTrace.create(o._diff(sym), over_space=s)

    def _adjoint(self):
        # there is a rule Tr[A^\dagger] -> Tr[A]^\dagger, which we don't want
        # to counteract here with an inverse rule
        return Adjoint(self)

    def _pseudo_inverse(self):
        return PseudoInverse(self)

class Adjoint(QuantumAdjoint, Operator):
    """Symbolic Adjoint of an operator"""

    simplifications = [scalars_to_op, delegate_to_method('_adjoint')]
    # The reason that Adjoint does not have have `match_replace` in
    # `simplifications`, respectively a `_rules` class attribute is that the
    # `_adjoint` property that we delegate to is mandatory. Thus, if we had
    # rules on top of that, it would create the confusing situation of the rule
    # contradicting the `_adjoint` property.

    def _pseudo_inverse(self):
        return self.operand.pseudo_inverse().adjoint()

class OperatorPlusMinusCC(SingleQuantumOperation, Operator):
    """An operator plus or minus its complex conjugate."""

    def __init__(self, op, *, sign=+1):
        self._sign = sign
        super().__init__(op, sign=sign)

    @property
    def kwargs(self):
        if self._sign > 0:
            return {
                'sign': +1,
            }
        else:
            return {
                'sign': -1,
            }

    @property
    def minimal_kwargs(self):
        if self._sign == +1:
            return {}
        else:
            return self.kwargs

    def _expand(self):
        return self

    def _diff(self, sym):
        return OperatorPlusMinusCC(self.operands._diff(sym), sign=self._sign)

    def _adjoint(self):
        return OperatorPlusMinusCC(self.operand.adjoint(), sign=self._sign)

    def _pseudo_inverse(self):
        return PseudoInverse(self.doit())

    def doit(self, classes=None, recursive=True, **kwargs):
        """Write out the complex conjugate summand

        See :meth:`.Expression.doit`.
        """
        return super().doit(classes, recursive, **kwargs)

    def _doit(self, **kwargs):
        if self._sign > 0:
            return self.operand + self.operand.adjoint()
        else:
            return self.operand - self.operand.adjoint()

class PseudoInverse(SingleQuantumOperation, Operator):
    r"""Unevaluated pseudo-inverse $\Op{X}^+$ of an operator $\Op{X}$.

    It is defined via the relationship

    .. math::

        \Op{X} \Op{X}^+ \Op{X} =  \Op{X}  \\
        \Op{X}^+ \Op{X} \Op{X}^+ =  \Op{X}^+  \\
        (\Op{X}^+ \Op{X})^\dagger = \Op{X}^+ \Op{X}  \\
        (\Op{X} \Op{X}^+)^\dagger = \Op{X} \Op{X}^+
    """
    simplifications = [scalars_to_op, delegate_to_method('_pseudo_inverse')]
    # `PseudoInverse` does not use rules because it delegates to
    # `_pseudo_inverse`, cf. `Adjoint`

    def _expand(self):
        return self

    def _pseudo_inverse(self):
        return self.operand

    def _adjoint(self):
        return Adjoint(self)

class NullSpaceProjector(SingleQuantumOperation, Operator):
    r"""Projection operator onto the nullspace of its operand.

    Returns the operator :math:`\mathcal{P}_{{\rm Ker} X}` with

    .. math::

        X \mathcal{P}_{{\rm Ker} X}
          = 0
        \Leftrightarrow
        X (1 - \mathcal{P}_{{\rm Ker} X})
          = X \\
        \mathcal{P}_{{\rm Ker} X}^\dagger
          = \mathcal{P}_{{\rm Ker} X}
          = \mathcal{P}_{{\rm Ker} X}^2
    """

    _rules = OrderedDict()
    simplifications = [
        scalars_to_op,
        match_replace,
    ]

    def _expand(self):
        return self

    def _adjoint(self):
        return Adjoint(self)

    def _pseudo_inverse(self):
        return PseudoInverse(self)

class OperatorIndexedSum(QuantumIndexedSum, Operator):
    """Indexed sum over operators"""

    _rules = OrderedDict()
    simplifications = [
        assoc_indexed,
        scalars_to_op,
        indexed_sum_over_kronecker,
        indexed_sum_over_const,
        match_replace,
    ]

    def _pseudo_inverse(self):
        return PseudoInverse(self)

###############################################################################
# Constructor Routines
###############################################################################

tr = OperatorTrace.create

###############################################################################
# Auxilliary routines
###############################################################################

def factor_for_trace(ls: HilbertSpace, op: Operator) -> Operator:
    r"""Factor `ls` out of `op` for easy tracing.

    Given a :class:`.LocalSpace` `ls` to take the partial trace over and an
    operator `op`, factor the trace such that operators acting on disjoint
    degrees of freedom are pulled out of the trace. If the operator acts
    trivially on ls the trace yields only a pre-factor equal to the dimension
    of ls. If there are :class:`.LocalSigma` operators among a product, the
    trace's cyclical property is used to move to sandwich the full product by
    :class:`LocalSigma` operators:

    .. math::

        {\rm Tr} A \sigma_{jk} B = {\rm Tr} \sigma_{jk} B A \sigma_{jj}

    Args:
        ls: Degree of freedom to trace over
        op: Operator to take the trace of

    Returns:
        The (partial) trace over the operator's spc-degrees of freedom
    """
    if op.space == ls:
        if isinstance(op, OperatorTimes):
            pull_out = [o for o in op.operands if o.space is TrivialSpace]
            rest = [o for o in op.operands if o.space is not TrivialSpace]
            if pull_out:
                return OperatorTimes.create(*pull_out) * OperatorTrace.create(
                    OperatorTimes.create(*rest), over_space=ls
                )
        raise CannotSimplify()
    if ls & op.space == TrivialSpace:
        return ls.dimension * op
    if ls < op.space and isinstance(op, OperatorTimes):
        pull_out = [o for o in op.operands if (o.space & ls) == TrivialSpace]

        rest = [o for o in op.operands if (o.space & ls) != TrivialSpace]
        if not isinstance(rest[0], LocalSigma) or not isinstance(
            rest[-1], LocalSigma
        ):
            for j, r in enumerate(rest):
                if isinstance(r, LocalSigma):
                    m = r.j
                    rest = (
                        rest[j:] + rest[:j] + [LocalSigma.create(m, m, hs=ls)]
                    )
                    break
        if not rest:
            rest = [IdentityOperator]
        if len(pull_out):
            return OperatorTimes.create(*pull_out) * OperatorTrace.create(
                OperatorTimes.create(*rest), over_space=ls
            )
    raise CannotSimplify()

def decompose_space(H, A):
    """Simplifies OperatorTrace expressions over tensor-product spaces by
    turning it into iterated partial traces.

    Args:
        H (ProductSpace): The full space.
        A (Operator):

    Returns:
        Operator: Iterative partial trace expression
    """
    return OperatorTrace.create(
        OperatorTrace.create(A, over_space=H.operands[-1]),
        over_space=ProductSpace.create(*H.operands[:-1]),
    )

def get_coeffs(expr, expand=False, epsilon=0.0):
    """Create a dictionary with all Operator terms of the expression
    (understood as a sum) as keys and their coefficients as values.

    The returned object is a defaultdict that return 0. if a term/key
    doesn't exist.

    Args:
        expr: The operator expression to get all coefficients from.
        expand: Whether to expand the expression distributively.
        epsilon: If non-zero, drop all Operators with coefficients that have
            absolute value less than epsilon.

    Returns:
        dict: A dictionary ``{op1: coeff1, op2: coeff2, ...}``
    """
    if expand:
        expr = expr.expand()
    ret = defaultdict(int)
    operands = expr.operands if isinstance(expr, OperatorPlus) else [expr]
    for e in operands:
        c, t = _coeff_term(e)
        try:
            if abs(complex(c)) < epsilon:
                continue
        except TypeError:
            pass
        ret[t] += c
    return ret

def _coeff_term(op):
    # TODO: remove
    if isinstance(op, ScalarTimesOperator):
        return op.coeff, op.term
    elif is_scalar(op):
        if op == 0:
            return 0, ZeroOperator
        else:
            return op, IdentityOperator
    else:
        return 1, op

def factor_coeff(cls, ops, kwargs):
    """Factor out coefficients of all factors."""
    coeffs, nops = zip(*map(_coeff_term, ops))
    coeff = 1
    for c in coeffs:
        coeff *= c
    if coeff == 1:
        return nops, coeffs
    else:
        return coeff * cls.create(*nops, **kwargs)

def adjoint(obj):
    """Return the adjoint of an obj."""
    try:
        return obj.adjoint()
    except AttributeError:
        return obj.conjugate()

###############################################################################
# Extra ("manual") simplifications
###############################################################################

def rewrite_with_operator_pm_cc(expr):
    """Try to rewrite expr using :class:`.OperatorPlusMinusCC`.

    Example::

        >>> A = OperatorSymbol('A', hs=1)
        >>> sum = A + A.dag()
        >>> sum2 = rewrite_with_operator_pm_cc(sum)
        >>> print(ascii(sum2))
        A^(1) + c.c.
    """
    # TODO: move this to the toolbox
    from qalgebra.toolbox.core import temporary_rules

    def _combine_operator_p_cc(A, B):
        if B.adjoint() == A:
            return OperatorPlusMinusCC(A, sign=+1)
        else:
            raise CannotSimplify

    def _combine_operator_m_cc(A, B):
        if B.adjoint() == A:
            return OperatorPlusMinusCC(A, sign=-1)
        else:
            raise CannotSimplify

    def _scal_combine_operator_pm_cc(c, A, d, B):
        if B.adjoint() == A:
            if c == d:
                return c * OperatorPlusMinusCC(A, sign=+1)
            elif c == -d:
                return c * OperatorPlusMinusCC(A, sign=-1)
        raise CannotSimplify

    A = wc("A", head=Operator)
    B = wc("B", head=Operator)
    c = wc("c", head=Scalar)
    d = wc("d", head=Scalar)

    with temporary_rules(OperatorPlus, clear=True):
        OperatorPlus.add_rule(
            'PM1', pattern_head(A, B), _combine_operator_p_cc
        )
        OperatorPlus.add_rule(
            'PM2',
            pattern_head(pattern(ScalarTimesOperator, -1, B), A),
            _combine_operator_m_cc,
        )
        OperatorPlus.add_rule(
            'PM3',
            pattern_head(
                pattern(ScalarTimesOperator, c, A),
                pattern(ScalarTimesOperator, d, B),
            ),
            _scal_combine_operator_pm_cc,
        )
        return expr.rebuild()

Operator._zero = ZeroOperator
Operator._one = IdentityOperator
Operator._base_cls = Operator
Operator._scalar_times_expr_cls = ScalarTimesOperator
Operator._plus_cls = OperatorPlus
Operator._times_cls = OperatorTimes
Operator._adjoint_cls = Adjoint
Operator._indexed_sum_cls = OperatorIndexedSum
Operator._derivative_cls = OperatorDerivative