text:
'''
Provides methods for plotting data as spatial and timeseries plots
'''
import matplotlib
import matplotlib.pyplot as plt
import numpy as _np
import pandas

import awrams.utils.datetools as datetools
import awrams.utils.extents as extents
from awrams.utils.metatypes import ObjectDict as _od
from awrams.utils.metatypes import DataDict as _dd
from awrams.visualisation import layout as _layout

from .support import _sanitize_inputs,_sanitize_period,_sanitize_aggregate_method
from .support import spatial_aggregate as _spatial_aggregate

ROW_FIELD = 'source'
COLUMN_FIELD = 'variable'

def _plot_grids(grids):
    '''
    Plot the results of a query_grids() operation
    '''
    layout = _layout.DefaultSpatialGridLayout('period','variable')
    gridview = layout.generate_view(grids)
    gridview.draw()

def _plot_spatial(source,pattern=None):
    '''
    Plot spatial (gridded) data according to supplied
    pattern
    '''
    if pattern is None:
        pattern = {}
    grids = source.query_grids(pattern)
    _plot_grids(grids)

def _plot_nans(source,pattern):
    '''
    Show locations of NaN values in data
    '''
    grids = source.query_grids(pattern)
    for datum in grids.results:
        datum.data = _np.isnan(datum.data)
    layout = _layout.ShowNaNGridLayout('period','variable')
    gridview = layout.generate_view(grids)
    gridview.draw()

def _plot_filter(source,filter_fn,pattern = None):
    '''
    Plot results filtered by binary filter_fn
    e.g filter_fn = numpy.isnan
    Will highlight NaN values
    '''
    if pattern is None:
        pattern = {}
    grids = source.query_grids(pattern)
    for datum in grids.results:
        datum.data = filter_fn(datum.data)
    layout = _layout.ShowNaNGridLayout('period','variable')
    gridview = layout.generate_view(grids)
    gridview.draw()

def show_extent(extent,context):
    '''
    Show the region covered by extent, plotted against
    the supplied context (default to continental)

    :param extent: extent of interest
    :param context: extent to plot relative to, defaults to extent.default()
    :return: None
    '''

    c_ext = extent.translate(context.geospatial_reference())

    to_plot = _od()
    to_plot.extent = context
    to_plot.data = _np.ma.ones(context.shape)
    to_plot.data.mask = context.mask
    to_plot.data[c_ext.indices][c_ext.mask==False] = 2.0

    gridview = _layout.GridView(1,1)

    cmap = _layout.highlight_map((0.1,0.3,0.1),(1.0,0,0),(0.1,0.4,0.8))

    gridview[0,0] = _layout.NoBarView(to_plot,{'title': extent.__repr__(),'cmap': cmap})

    gridview.draw()

def _default_labels_single_grid(from_dict):
    labels = _od()
    labels.title = "%s, (%s)" % ( from_dict['variable'].name, from_dict['variable'].source.name )
    labels.ylabel = "%s" % (datetools.pretty_print_period(from_dict['period']))
    labels.cmap = 'RdYlBu'
    labels.units = from_dict['variable'].units
    return labels

def _trunc(x, n):
    '''Truncates/pads a float x to n decimal places without rounding'''
    slen = len('%.*f' % (n, x))
    return str(x)[:slen]

def _plot_spatial_multi(var_group,period,extent=None,**kwds):
    if not extent:
        extent = var_group.source.extent

    vg_dict = _dd()

    for v in var_group:
        aggregator = _sanitize_aggregate_method(v)

        data = aggregator(v,period,extent)

        vg_dict.add_query_item([ROW_FIELD,COLUMN_FIELD],_od(period=period,_data=data,units=v.units,extent=extent,variable=v.name,source=v.source))

    layout = _layout.DefaultSpatialGridLayout(ROW_FIELD,COLUMN_FIELD,**kwds)

    gridview = layout.generate_view(vg_dict)
    gridview.draw()
    return gridview

def get_range(inputs,extent=None,periods=None,colapse_dimension='time'):
    '''
    Return the minimum and maximum values across all periods and extents and variables
    For use when specifying a colorbar range or y-axis limits

    :param inputs: ResultVariable or VariableGroup
    :param extent: extents object
    :param periods: pandas DatetimeIndex or list of pandas DatetimeIndex's
    :param colapse_dimension: 'time' or 'spatial'
    :return: minimum,maximum
    '''

    if not type(periods) in (str,tuple):
        periods = [periods]

    abs_min = 1e10
    abs_max = -1e10
    for singleperiod in periods:

        inputs = _sanitize_inputs(inputs)

        period = _sanitize_period(singleperiod,inputs.source.period.freq)

        if not extent:
            extent = inputs.source.extent

        for v in inputs:
            if colapse_dimension == 'time':
                aggregator = _sanitize_aggregate_method(v)
            else:
                aggregator = _spatial_aggregate

            data = aggregator(v, period, extent)

            local_min = data.min()
            local_max = data.max()
            if local_min < abs_min:
                abs_min = local_min
            if local_max > abs_max:
                abs_max = local_max

    abs_min = float(_trunc(round(abs_min, 1), 1))
    abs_max = float(_trunc(abs_max + 0.1, 1))

    return abs_min, abs_max

def spatial(inputs,period,extent=None,**kwds):
    '''
    For each period calculate the minimum and max values, then
    use these extremes for consistent colorbar values,
    If the minval and maxval values are passed in, use those instead
    Use this to get plot colour consistency across data
    indicated in multiple calls
    '''

    inputs = _sanitize_inputs(inputs)

    period = _sanitize_period(period,inputs.source.period.freq)

    if not extent:
        extent = inputs.source.extent

    return _plot_spatial_multi(inputs,period,extent,**kwds)

def timeseries(inputs,period=None,extent=None,**kwds):
    '''
    Inputs is a variable or group of variables
    If extent is larger than a single cell, variables
    will be aggregated using their default method
    '''

    inputs = _sanitize_inputs(inputs)

    if period is None:
        period = inputs.source.period
    else:
        period = _sanitize_period(period, inputs.source.period.freq)

    if extent is None:
        extent = inputs.source.extent

    vg_dict = _dd()

    for v in inputs:
        data = _spatial_aggregate(v,period,extent)
        vg_dict.add_item(_od(period=period,_data=data,units=v.units,extent=extent,method='mean',variable=v.name,source=v.source))

    layout = _layout.SpatialAggregateLayout(**kwds)

    gridview = layout.generate_view(vg_dict)
    gridview.draw()
    return gridview


text:
import numpy as np
import re
from . import stereonet_math as smath

def parse_strike_dip(strike, dip):
    """
    Parses strings of strike and dip and returns strike and dip measurements
    following the right-hand-rule.

    Dip directions are parsed, and if the measurement does not follow the
    right-hand-rule, the opposite end of the strike measurement is returned.

    Accepts either quadrant-formatted or azimuth-formatted strikes.

    For example, this would convert a strike of "N30E" and a dip of "45NW" to
    a strike of 210 and a dip of 45.

    Parameters
    ----------
    strike : string
        A strike measurement. May be in azimuth or quadrant format.
    dip : string
        The dip angle and direction of a plane.

    Returns
    -------
    azi : float
        Azimuth in degrees of the strike of the plane with dip direction
        indicated following the right-hand-rule.
    dip : float
        Dip of the plane in degrees.
    """
    strike = parse_azimuth(strike)
    dip, direction = split_trailing_letters(dip)

    if direction is not None:
        expected_direc = strike + 90
        if opposite_end(expected_direc, direction):
            strike += 180

    if strike > 360:
        strike -= 360

    return strike, dip

def parse_rake(strike, dip, rake):
    """
    Parses strings of strike, dip, and rake and returns a strike, dip, and rake
    measurement following the right-hand-rule, with the "end" of the strike
    that the rake is measured from indicated by the sign of the rake (positive
    rakes correspond to the strike direction, negative rakes correspond to the
    opposite end).

    Accepts either quadrant-formatted or azimuth-formatted strikes.

    For example, this would convert a strike of "N30E", dip of "45NW", with a
    rake of "10NE" to a strike of 210, dip of 45, and rake of 170.

    Rake angles returned by this function will always be between 0 and 180

    If no directions are specified, the measuriement is assumed to follow the
    usual right-hand-rule convention.

    Parameters
    ----------
    strike : string
        A strike measurement. May be in azimuth or quadrant format.
    dip : string
        The dip angle and direction of a plane.
    rake : string
        The rake angle and direction that the rake is measured from.

    Returns
    -------
    strike, dip, rake : floats
        Measurements of strike, dip, and rake following the conventions
        outlined above.
    """
    strike, dip = parse_strike_dip(strike, dip)
    rake, direction = split_trailing_letters(rake)

    if direction is not None:
        if opposite_end(strike, direction):
            rake = -rake

    if rake < 0:
        rake += 180
    elif rake > 180:
        rake -= 180

    return strike, dip, rake

def parse_plunge_bearing(plunge, bearing):
    """
    Parses strings of plunge and bearing and returns a consistent plunge and
    bearing measurement as floats. Plunge angles returned by this function will
    always be between 0 and 90.

    If no direction letter(s) is present, the plunge is assumed to be measured
    from the end specified by the bearing. If a direction letter(s) is present,
    the bearing will be switched to the opposite (180 degrees) end if the
    specified direction corresponds to the opposite end specified by the
    bearing.

    Parameters
    ----------
    plunge : string
        A plunge measurement.
    bearing : string
        A bearing measurement. May be in azimuth or quadrant format.

    Returns
    -------
    plunge, bearing: floats
        The plunge and bearing following the conventions outlined above.

    Examples
    ---------

        >>> parse_plunge_bearing("30NW", 160)
        ... (30, 340)

    """
    bearing = parse_azimuth(bearing)
    plunge, direction = split_trailing_letters(plunge)

    if direction is not None:
        if opposite_end(bearing, direction):
            bearing +=180

    if plunge < 0:
        bearing += 180
        plunge = -plunge
    if plunge > 90:
        bearing += 180
        plunge = 180 - plunge

    if bearing > 360:
        bearing -= 360

    return plunge, bearing

def dip_direction2strike(azimuth):
    """
    Converts a planar measurment of dip direction using the dip-azimuth
    convention into a strike using the right-hand-rule.

    Parameters
    ----------
    azimuth : number or string
        The dip direction of the plane in degrees. This can be either a
        numerical azimuth in the 0-360 range or a string representing a quadrant
        measurement (e.g. N30W).

    Returns
    -------
    strike : number
        The strike of the plane in degrees following the right-hand-rule.
    """
    azimuth = parse_azimuth(azimuth)
    strike = azimuth - 90
    if strike < 0:
        strike += 360
    return strike

def strike2dip_direction(strike):
    """
    Converts a planar measurement of strike using the right-hand-rule into the
    dip direction (i.e. the direction that the plane dips).

    Parameters
    ----------
    strike : number or string
        The strike direction of the plane in degrees. This can be either a
        numerical azimuth in the 0-360 range or a string representing a quadrant
        measurement (e.g. N30W).

    Returns
    -------
    azimuth : number
        The dip direction of the plane in degrees (0-360).
    """
    strike = parse_azimuth(strike)
    dip_direction = strike + 90
    if dip_direction > 360:
        dip_direction -= 360
    return dip_direction

def opposite_end(azimuth, quadrant):
    direc = quadrantletter_to_azimuth(quadrant)
    dot = np.dot(_azimuth2vec(direc), _azimuth2vec(azimuth))
    return dot < 0

def split_trailing_letters(item):
    # Match things like 45NW, 45E, 45WNW, etc
    letters = re.search(r'[NESWnesw]', item)
    if letters is None:
        return float(item), None
    else:
        measurement = item[:letters.start()]
        direction = item[letters.start():]
        return float(measurement), direction

def circmean(azimuths):
    azimuths = np.radians(azimuths)
    x = np.cos(azimuths)
    y = np.sin(azimuths)
    return np.degrees(np.arctan2(y.mean(), x.mean()))

def quadrantletter_to_azimuth(letters):
    letters = letters.strip()
    azimuth = {'N':0, 'S':180, 'E':90, 'W':270}

    azi = azimuth[letters[-1]]
    for letter in letters[1::-1]:
        azi = circmean([azi, azimuth[letter]])
    return azi

def _azimuth2vec(azi):
    azi = np.radians(azi)
    return np.cos(azi), np.sin(azi)

def parse_azimuth(azimuth):
    """
    Parses an azimuth measurement in azimuth or quadrant format.

    Parameters
    -----------
    azimuth : string or number
        An azimuth measurement in degrees or a quadrant measurement of azimuth.

    Returns:
    --------
    azi : float
        The azimuth in degrees clockwise from north (range: 0-360)

    See Also:
    ---------
    `parse_quadrant_measurement`
    `parse_strike_dip`
    `parse_plunge_bearing`
    """
    try:
        azimuth = float(azimuth)
    except ValueError:
        if not azimuth[0].isalpha():
            raise ValueError('Ambiguous azimuth: {}'.format(azimuth))
        azimuth = parse_quadrant_measurement(azimuth)
    return azimuth

def parse_quadrant_measurement(quad_azimuth):
    """
    Parses a quadrant measurement of the form "AxxB", where A and B are cardinal
    directions and xx is an angle measured relative to those directions.

    In other words, it converts a measurement such as E30N into an azimuth of
    60 degrees, or W10S into an azimuth of 260 degrees.

    For ambiguous quadrant measurements such as "N30S", a ValueError is raised.

    Parameters
    -----------
    quad_azimuth : string
        An azimuth measurement in quadrant form.

    Returns:
    --------
    azi : float
        An azimuth in degrees clockwise from north.

    See Also:
    ---------
    `parse_azimuth`
    """
    def rotation_direction(first, second):
        return np.cross(_azimuth2vec(first), _azimuth2vec(second))

    # Parse measurement
    quad_azimuth = quad_azimuth.strip()
    try:
        first_dir = quadrantletter_to_azimuth(quad_azimuth[0].upper())
        sec_dir = quadrantletter_to_azimuth(quad_azimuth[-1].upper())
    except KeyError:
        raise ValueError('{} is not a valid azimuth'.format(quad_azimuth))

    angle = float(quad_azimuth[1:-1])

    # Convert quadrant measurement into an azimuth
    direc = rotation_direction(first_dir, sec_dir)
    azi = first_dir + direc * angle

    # Catch ambiguous measurements such as N10S and raise an error
    if abs(direc) < 0.9:
        raise ValueError('{} is not a valid azimuth'.format(quad_azimuth))

    # Ensure that 0 <= azi <= 360
    if azi < 0:
        azi += 360
    elif azi > 360:
        azi -= 360

    return azi

ext: py

max_stars_repo_name: spiessbuerger/mplstereonet

max_issues_repo_name: spiessbuerger/mplstereonet

max_forks_repo_name: spiessbuerger/mplstereonet

reason: import numpy


text:
# %% Imports
import numpy as np
import cppimport.import_hook
sys.path.insert(0, "src/tallem")

# %% Random data
X = np.random.uniform(size=(10, 2))
y = np.random.uniform(size=(10, 1))

# %% Neighborhood graph
import numpy as np
X = np.random.uniform(size=(10, 2))
Y = np.random.uniform(size=(5, 2))

## Get the k-nearest neighbors in X around the points in Y
from scipy.spatial import KDTree
from scipy.sparse import csc_matrix, lil_matrix
import os

G = lil_matrix((X.shape[0], Y.shape[0]))

## X, Y point sets, distances are minkowski, knn-graph desired
tree = KDTree(X)
knn_in_x = tree.query(Y, k = 5, workers=os.cpu_count())
knn_dist, knn_ind = knn_in_x
for j in range(Y.shape[0]):
	G[knn_ind[j,:],j] = knn_dist[j,:]
# G.tocsc()

## X, Y point sets, distances are minkowski, ball-graph desired
from src.tallem.distance import dist
tree = KDTree(X)
eps_in_x = tree.query_ball_point(Y, r=0.35, workers=os.cpu_count())
for j, ind in enumerate(eps_in_x):
	if len(ind) > 0:
		G[ind,j] = dist(Y[j,:], X[ind,:])
		# dist(X[np.array(ind),:], Y[[j],:])

# %% Test all the SVD implementations
bn = example.BetaNuclearDense(1000, 3, 3)
bn.three_svd()

# %% 2 x 2 SVD
# Based on: https://lucidar.me/en/mathematics/singular-value-decomposition-of-a-2x2-matrix/
def svd_2x2(x):
	a,b,c,d = x[0,0], x[0,1], x[1,0], x[1,1]
	theta = 0.5*np.arctan2(2*a*c + 2*b*d, a**2 + b**2 - c**2 - d**2)
	phi = 0.5*np.arctan2(2*a*b + 2*c*d, a**2 - b**2 + c**2 - d**2)
	ct,st,cp,sp = np.cos(theta),np.sin(theta),np.cos(phi),np.sin(phi)
	u = np.array([[ct, -st],[st, ct]], dtype=np.float32)
	s1 = a**2 + b**2 + c**2 + d**2
	s2 = np.sqrt((a**2 + b**2 - c**2 - d**2)**2 + 4*(a*c + b*d)**2)
	s = np.array([ np.sqrt((s1 + s2)/2), np.sqrt((s1 - s2)/2) ])
	sign11 = 1 if ((a*ct+c*st)*cp + (b*ct + d*st)*sp) >= 0 else -1
	sign22 = 1 if ((a*st-c*ct)*sp + (-b*st + d*ct)*cp) >= 0 else -1
	v = np.array([[sign11*cp,-sign22*sp],[sign11*sp,sign22*cp]], dtype=np.float32)
	return((u,s,v))

y = np.hstack((x, np.zeros(shape=(3,1))))
print(np.linalg.svd(x))
print(np.linalg.svd(y))

# %% 2 x 2 SVD
import numpy as np
from tallem import fast_svd
x = np.random.uniform(size=(3,3))
y = fast_svd.fast_svd(x)

# %% Try import hook
import numpy as np
x = np.random.uniform(size=(3,3))
import cppimport.import_hook
sys.path.insert(0, "src/tallem")
import fast_svd

fast_svd.fast_svd(x)

# %% carma
import numpy as np
from tallem import carma_svd
X = np.random.uniform(size=(10,10))
y = np.random.uniform(size=(10,1))
print(carma_svd.ols(X, y))

# %% stiefel loss
import numpy as np
from tallem import fast_svd
d, D, J, n = 3, 3, 500, 10000
A = np.random.uniform(size=(d*J, D))
Phi = np.random.uniform(size=(d*J,d*n))

[A, R] = np.linalg.qr(A)

stf = fast_svd.StiefelLoss(n, d, D)
stf.output = A.T @ Phi

# %% stiefel gradient
%%time
f, gf = stf.gradient(Phi)

# %% stiefel gradient (3 x 3)
%%time
stf.three_svd()

# %% stiefel gradient carma (3 x 3)
%%time
stf.three_svd_carma()

# %% base python + numpy
%%time
nuclear_norm = 0.0
gr = np.zeros((D, d))
for i in range(n):
	u,s,vt = np.linalg.svd(stf.output[:,(i*3):(i*3+3)], compute_uv=True, full_matrices=False)
	nuclear_norm += np.sum(np.abs(s))
	gr += u @ vt

# %% stiefel gradient
## Compare with python loop

# %% testing tallem
import numpy as np
from tallem.datasets import mobius_band
from tallem.cover import IntervalCover
from tallem.cover import partition_of_unity

M, f = mobius_band(n_polar = 26, n_wide = 6, scale_band = 0.1, embed = 3).values()
B = f[:,1:2]
cover = IntervalCover(B, n_sets = 10, overlap = 0.75, gluing=[1])
PoU = partition_of_unity(B, cover, beta="triangular")

# %% Plot data set
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
plt.axis('scaled')
ax.scatter(M[:,0], M[:,1], M[:,2], marker='o', c=f[:,1])

# %% IPyVolume plot
# import ipyvolume as ipv
# import matplotlib.cm as cm
# import bqplot.scales
# scales = {
#     'x': bqplot.scales.LogScale(min=10**-3, max=10**3),
#     'y': bqplot.scales.LinearScale(min=-3, max=3),
#     'z': bqplot.scales.LinearScale(min=-3, max=3),
# }
# color_scale = bqplot.scales.ColorScale(min=0, max=3, colors=["#f00", "#0f0", "#00f"])
# cm.viridis(B).
# fig = ipv.figure()
# sp = ipv.scatter(M[:,0], M[:,1], M[:,2], size=5, marker="sphere", color=cm.viridis(B))
# ipv.show()

# %% Local models
from tallem.mds import classical_MDS
from tallem.distance import dist
local_dim = 2
local_map = lambda x: classical_MDS(dist(x, as_matrix=True), k = local_dim)
local_models = { index : local_map(M[subset,:]) for index, subset in cover }

## Rotation, scaling, translation, and distance information for each intersecting cover subset
from tallem.procrustes import align_models
alignments = align_models(cover, local_models)

# %% Frame reduction (optimization)
from tallem.stiefel import frame_reduction
D = 3 # chosen dimension
A, stf = frame_reduction(alignments, PoU, D, fast_gradient=True)

# %% Run the assembly
from tallem.procrustes import global_translations

## Construct the global assembly function
assembly = np.zeros((stf.n, D), dtype=np.float64)
coords = np.zeros((1,D), dtype=np.float64)
index_set = list(local_models.keys())
translations = global_translations(alignments)
for i in range(stf.n):
	w_i = np.ravel(PoU[i,:].todense())
	nz_ind = np.where(w_i > 0)[0]
	d_frame = stf.get_frame(i) # already weighted
	coords.fill(0)
	## Construct assembly functions F_j(x) for x_i
	for j in nz_ind:
		subset_j = cover[index_set[j]]
		relative_index = np.searchsorted(subset_j, i)
		u, s, vt = np.linalg.svd(A @ (A.T @ stf.generate_frame(j, w_i)), full_matrices=False, compute_uv=True)
		d_coords = local_models[index_set[j]][relative_index,:]
		coords += w_i[j]*A.T @ (u @ vt) @ (d_coords + translations[j])
	assembly[i,:] = coords

# %% Plot embedding
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(assembly[:,0], assembly[:,1], assembly[:,2], marker='o', c=f[:,0])

# %% Plot embedding 2
from tallem import tallem_transform
assembly2 = tallem_transform(M, B)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(assembly2[:,0], assembly2[:,1], assembly2[:,2], marker='o', c=f[:,0])

# %% Reference old method
ref_M = tallem_transform(M, B)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(ref_M[:,0], ref_M[:,1], ref_M[:,2], marker='o', c=f[:,0])

# %% Initialize Frame stuff
from tallem import fast_svd
d, D, J, n = local_dim, 3, len(cover), cover.n_points
stf = fast_svd.StiefelLoss(n, d, D)

## Initialize rotation matrix hashmap
I1 = [index[0] for index in alignments.keys()]
I2 = [index[1] for index in alignments.keys()]
R = np.vstack([pa['rotation'] for index, pa in alignments.items()])
stf.init_rotations(I1, I2, R, J)

## Populate frame matrix map
for i in range(n):
	stf.populate_frame(i, np.ravel(PoU[i,:].todense()), False)

# %% Extract the frames
Fb = stf.all_frames()
Eval, Evec = np.linalg.eigh(Fb @ Fb.T)
A0 = Evec[:,np.argsort(-Eval)[:D]]

def stiefel_gradient(d_frame):
	stf.embed(d_frame.T) ## populates (D x dn) output array
	F, GF = stf.gradient()
	return(F / stf.n, GF / stf.n)

F, GF = stiefel_gradient(A0, stf)

for delta in range(10):
	loss, _ = stiefel_gradient(A0 - (delta/10)*GF, stf)
	print(loss)

from pymanopt.manifolds import Stiefel

Phi = stf.all_frames()
import autograd.numpy as auto_np
def sv_loss(A):
	out = A.T @ Phi
	val = 0.0
	for i in range(n):
		u,s,v = auto_np.linalg.svd(out[:,i*d:((i+1)*d)], full_matrices=False)
		val += auto_np.sum(auto_np.abs(s))
	return(-val/n)

from autograd import grad
gr = grad(sv_loss)
gr(A0)

from pymanopt.manifolds import Stiefel
from pymanopt import Problem
from pymanopt.solvers import SteepestDescent
manifold = Stiefel(d*J, d)
problem = Problem(manifold=manifold, cost=sv_loss)
solver = SteepestDescent()
Xopt = solver.solve(problem, x=A0)

manifold.retr(X=A0, G=G.T)

def stiefel_loss(A):
	stf.embed(A.T) ## populates (D x dn) output array
	F, GF = stf.gradient()
	return(F / stf.n)
def stiefel_egrad(A):
	stf.embed(A.T) ## populates (D x dn) output array
	F, GF = stf.gradient()
	return(GF / stf.n)
manifold = Stiefel(d*J, d)
problem = Problem(manifold=manifold, cost=stiefel_loss, egrad=stiefel_egrad)
solver = SteepestDescent(mingradnorm=1e-12, maxiter=100, minstepsize=1e-14)
Xopt = solver.solve(problem, x=A0)

#-8.4632914625038969e-01	9.29798595e-07

from tallem.diagnostics import check_gradient
check_gradient(problem, x=A0)

# %% Procrustes
import matplotlib.pyplot as plt
import numpy as np
star1 = np.array([[131,38], [303, 39], [357, 204], [217,305], [78,204]])
star2 = np.array([[61, 33],[120, 73],[103,162],[32,177],[7,97]])
plt.scatter(star1[:,0], star1[:,1], c="red")
plt.scatter(star2[:,0], star2[:,1], c="blue")

from tallem.procrustes import opa, old_procrustes
pa = old_procrustes(star1, star2)

plt.scatter(A[:,0], A[:,1])
plt.scatter(B[:,0], B[:,1])

r,s,t,d = old_procrustes(star1, star2).values()

plt.scatter(star1[:,0], star1[:,1], c="red")
#Star2 = (pa['rotation'] @ star2.T).T + pa['translation']
Star2 = s * star2 @ r + t
plt.scatter(Star2[:,0], Star2[:,1], c="blue")

plt.scatter(star2[:,0], star2[:,1], c="blue")
#Star1 = ((star1 * 1/pa['scaling']) @ pa['rotation'].T) + pa['translation']
plt.scatter(Star1[:,0], Star1[:,1], c="red")

r,s,t,d = opa(a, b).values()
plt.scatter(a[:,0], a[:,1], c="red")
plt.scatter(b[:,0], b[:,1], c="blue")

Z = s*a@r+t
plt.scatter(Z[:,0], Z[:,1], c="orange")

#%% Debug
import numpy as np
from tallem.mds import classical_MDS
from tallem.distance import dist
from tallem.procrustes import align_models, global_translations
from tallem.stiefel import frame_reduction
from tallem.datasets import mobius_band
from tallem.cover import IntervalCover, partition_of_unity
from tallem.utility import find_where
from tallem.assembly import assemble_frames

d, D = 2, 3 # chosen dimensions
M, B = mobius_band(n_polar = 26, n_wide = 6, embed = 3).values()
f = B[:,1].reshape((B.shape[0], 1))
# M, f = X, F[:,1].reshape((X.shape[0], 1))
cover = IntervalCover(f, n_sets = 10, overlap = 0.30, gluing=[1])
local_map = lambda x: classical_MDS(dist(x, as_matrix=True), k = d)
local_models = { index : local_map(M[subset,:]) for index, subset in cover }

PoU = partition_of_unity(f, cover, beta="triangular")
alignments = align_models(cover, local_models)
A, stf = frame_reduction(alignments, PoU, D, fast_gradient=False, optimize = False)

translations = global_translations(alignments)
Y = assemble_frames(stf, A, cover, PoU, local_models, translations)

## Verify rotation matrices stored correctly
for i1,i2 in alignments.keys():
	diff = np.sum(np.abs(alignments[(i1,i2)]['rotation'] - stf.get_rotation(i1,i2,len(cover))))
	print(diff)
	# print(stf.get_rotation(i1,i2,len(cover)))

def phi_debug(point, cover_set = None):
	J = PoU.shape[1]
	k = np.argmax(PoU[point,:]) if cover_set is None else cover_set
	out = np.zeros((d*J, d))
	def weighted_omega(j):
		nonlocal point, k
		w = np.sqrt(PoU[point,j])
		pair_exists = np.array([pair in list(alignments.keys()) for pair in [(k,j), (j,k)]])
		if w == 0.0 or not(pair_exists.any()):
			return(w*np.eye(d))
		return(w*alignments[(k,j)]['rotation'] if pair_exists[0] else w*alignments[(j,k)]['rotation'].T)
	return(np.vstack([weighted_omega(j) for j in range(J)]))

### Verify Phi maps match
for i in range(cover.n_points):
	for j in range(len(cover)):
		p1 = phi_debug(point=i, cover_set=j)
		p2 = stf.generate_frame(j, np.ravel(np.sqrt(PoU[i,:]).todense()))
		diff = np.sum(np.abs(p1-p2))
		if (diff > 0.0):
			break

## Verify all_frames() is working
for i in range(cover.n_points):
	stf.populate_frame(i, np.sqrt(np.ravel(PoU[i,:].todense())), False)
Fb = np.hstack([phi_debug(point=i, cover_set=np.argmax(PoU[i,:])) for i in range(cover.n_points)])
np.sum(np.abs(Fb - stf.all_frames()))

## Verify initial frame guess is the same
Phi_N = Fb @ Fb.T # (dJ x dJ)
Eval, Evec = np.linalg.eigh(Phi_N)
a0_debug = Evec[:,np.argsort(-Eval)[:D]]
Fb = stf.all_frames()
Eval, Evec = np.linalg.eigh(Fb @ Fb.T)
a0_test = Evec[:,np.argsort(-Eval)[:D]]
np.sum(np.abs(a0_test - a0_debug))

# %% Test new interface
from tallem import TALLEM
from tallem.cover import IntervalCover
from tallem.cover import partition_of_unity

X, B = mobius_band(n_polar=26, n_wide=6, embed=3).values()
B_polar = B[:,1].reshape((B.shape[0], 1))
cover = IntervalCover(B_polar, n_sets = 10, overlap = 0.30, gluing=[1])
# cover = IntervalCover(B, n_sets = 10, overlap = 0.30, gluing=[1, 0])
f = lambda x: classical_MDS(dist(x, as_matrix=True), k = 2)

## Parameterize the TALLEM instance and transform the data
embedding = TALLEM(cover=cover, local_map=f, n_components=3)
X_transformed = embedding.fit_transform(X, B_polar)
# X_transformed = embedding.fit_transform(X, B)

## Draw a 3D projection
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
# ax.scatter(X_transformed[:,0], X_transformed[:,1], X_transformed[:,2], marker='o', c=B_polar)
ax.scatter(X_transformed[:,0], X_transformed[:,1], X_transformed[:,2], marker='o', c=B[:,0])

# %% Assembly
#cover = IntervalCover(f, n_sets = 10, overlap = 0.20, gluing=[1])
#PoU = partition_of_unity(f, cover, beta="triangular")

## PART OF IT IS THE ASSEMBLY
## Construct the global assembly function
assembly = np.zeros((stf.n, D), dtype=np.float64)
coords = np.zeros((1,D), dtype=np.float64)
index_set = list(local_models.keys())
translations = global_translations(alignments)
for i in range(stf.n):
	w_i = np.ravel(PoU[i,:].todense())
	nz_ind = np.where(w_i > 0)[0]
	d_frame = stf.get_frame(i) # already weighted
	coords.fill(0)
	## Construct assembly functions F_j(x) for x_i
	for j in nz_ind:
		subset_j = cover[index_set[j]]
		relative_index = find_where(i, subset_j, True) ## This should always be true!
		u, s, vt = np.linalg.svd(((A @ A.T) @ stf.generate_frame(j, w_i)), full_matrices=False, compute_uv=True)
		d_coords = local_models[index_set[j]][relative_index,:]
		coords += (w_i[j]*A.T @ (u @ vt) @ (d_coords + translations[j]).T).T
	assembly[i,:] = coords

import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
# ax.scatter(assembly[:,0], assembly[:,1], assembly[:,2], marker='o', c=f)
ax.scatter(assembly[:,0], assembly[:,1], assembly[:,2], marker='o', c=B[:,1])

# %% debug old
from tallem import tallem_transform
d, D = 2, 3 # chosen dimensions
# M, f = mobius_band(n_polar = 26, n_wide = 6, embed = 3).values()
M, f = X, F[:,1]
f = f.reshape((len(f), 1))
Y = tallem_transform(M, f)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(Y[:,0], Y[:,1], Y[:,2], marker='o', c=f,s=1)

# for angle in range(0, 360, 15):
# 	ax.view_init(30, angle)
# 	plt.draw()
# 	plt.pause(0.5)

max_stars_repo_name: peekxc/tallem

max_issues_repo_name: peekxc/tallem

max_forks_repo_name: peekxc/tallem

reason: import numpy,from scipy


text:
import os
import glob2
import time
import logging
from functools import wraps
import pickle
import re

import pandas as pd
import numpy as np

logger = logging.getLogger(__name__)
logger.setLevel("INFO")
handler = logging.StreamHandler()
logger.addHandler(handler)

## DEBUGGING

def timed(func):
    """This decorator prints the execution time for the decorated function."""

    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        logger.info("EXECUTION TIME : {} ran in {}s".format(func.__name__, round(end - start, 2)))
        return result

    return wrapper

## DATA MANIPULATION

def get_data(folder_path=None, file_name=None, file_path=None, index = 'md5', compression=None, verbose=True):
    '''
    Load pandas object from file path.

    Parameters:
    -----------
    file_path or ( folder_path, file_name ) : str
       Filepath/URL of pandas object. Specify one or the other.
    index : str, default 'md5'
       Column to use as the row labels of the dataframe.
    compression : {‘infer’, ‘gzip’, ‘bz2’, ‘zip’, ‘xz’, None}
       Decompression type. Set to None for no decompression.

    Returns:
    --------
    A pandas dataframe.
    '''

    #define file path
    if file_path is None:
        file_path = os.path.join(folder_path, file_name)

    #extract extension
    path, extension = os.path.splitext(file_path)
    if compression is not None:
        path, extension = os.path.splitext(path)

    #load data
    if extension == '.csv':
        df = pd.read_csv(file_path, index_col=index, compression=compression)
    if extension in ['.pickle', '.pkl']:
        df = pd.read_pickle(file_path, compression=compression)#.set_index(index)

    #print shape
    print('loaded data shape :'.ljust(21), df.shape)

    return df

## FILES MANIPULATION

def get_file_list(folder_path, extensions =['zip']):
    ''' Returns the list of files with all given extensions inside the specified folder in the current working directory, '''
    file_list = []
    if folder_path is not None:
        for extension in extensions:
            file_list.extend(glob2.glob(os.path.join(folder_path, '*.' + extension)))
    return file_list

def pickle_results(results, file, path='nested-cv-results'):
    file_path = os.path.join(path, file)
    pickle.dump(results, open( file_path, "wb" ), protocol=pickle.HIGHEST_PROTOCOL)

def unpickle_results(file_path=None, file=None, path='nested-cv-results'):
    if file_path is None:
        file_path = os.path.join(path, file)
    return pickle.load(open( file_path, "rb" ))


text:
import json
import os
import time
import numpy as np
import matplotlib.pyplot as plt
import logging
import re
from datetime import datetime

from qcodes.instrument.base import Instrument
from qcodes.utils import validators
from qcodes.instrument.parameter import ManualParameter

import zhinst.ziPython as zi

log = logging.getLogger(__name__)

##########################################################################
# Module level functions
##########################################################################

def gen_waveform_name(ch, cw):
    """
    Return a standard waveform name based on channel and codeword number.

    Note the use of 1-based indexing of the channels. To clarify, the
    'ch' argument to this function is 0-based, but the naming of the actual
    waveforms as well as the signal outputs of the instruments are 1-based.
    The function will map 'logical' channel 0 to physical channel 1, and so on.

    """
    return 'wave_ch{}_cw{:03}'.format(ch+1, cw)

def gen_partner_waveform_name(ch, cw):
    """
    Return a standard waveform name for the partner waveform of a dual-channel
    waveform. The physical channel indexing is 1-based where as the logical channel
    indexing (i.e. the argument to this function) is 0-based. To clarify, the
    'ch' argument to this function is 0-based, but the naming of the actual
    waveforms as well as the signal outputs of the instruments are 1-based.
    The function will map 'logical' channel 0 to physical channel 1, and so on.
    """
    return gen_waveform_name(2*(ch//2) + ((ch + 1) % 2), cw)

def merge_waveforms(chan0=None, chan1=None, marker=None):
    """
    Merges waveforms for channel 0, channel 1 and marker bits into a single
    numpy array suitable for being written to the instrument. Channel 1 and marker
    data is optional. Use named arguments to combine, e.g. channel 0 and marker data.
    """
    chan0_uint = None
    chan1_uint = None
    marker_uint = None

    # The 'array_format' variable is used internally in this function in order to
    # control the order and number of uint16 words that we put together for each
    # sample of the final array. The variable is essentially interpreted as a bit
    # mask where each bit indicates which channels/marker values to include in
    # the final array. Bit 0 for chan0 data, 1 for chan1 data and 2 for marker data.
    array_format = 0

    if chan0 is not None:
        chan0_uint = np.array((np.power(2, 15)-1)*chan0, dtype=np.uint16)
        array_format += 1
    if chan1 is not None:
        chan1_uint = np.array((np.power(2, 15)-1)*chan1, dtype=np.uint16)
        array_format += 2
    if marker is not None:
        marker_uint = np.array(marker, dtype=np.uint16)
        array_format += 4

    if array_format == 1:
        return chan0_uint
    elif array_format == 2:
        return chan1_uint
    elif array_format == 3:
        return np.vstack((chan0_uint, chan1_uint)).reshape((-2,), order='F')
    elif array_format == 4:
        return marker_uint
    elif array_format == 5:
        return np.vstack((chan0_uint, marker_uint)).reshape((-2,), order='F')
    elif array_format == 6:
        return np.vstack((chan1_uint, marker_uint)).reshape((-2,), order='F')
    elif array_format == 7:
        return np.vstack((chan0_uint, chan1_uint, marker_uint)).reshape((-2,), order='F')
    else:
        return []

def plot_timing_diagram(data, bits, line_length=30):
    """
    Takes list of 32-bit integer values as read from the 'raw/dios/0/data' device nodes and creates
    a timing diagram of the result. The timing diagram can be used for verifying that e.g. the
    strobe signal (A.K.A the toggle signal) is periodic.
    """
    def _plot_lines(ax, pos, *args, **kwargs):
        if ax == 'x':
            for p in pos:
                plt.axvline(p, *args, **kwargs)
        else:
            for p in pos:
                plt.axhline(p, *args, **kwargs)

    def _plot_timing_diagram(data, bits):
        plt.figure(figsize=(20, 0.5*len(bits)))

        t = np.arange(len(data))
        _plot_lines('y', 2*np.arange(len(bits)), color='.5', linewidth=2)
        _plot_lines('x', t[0:-1:2], color='.5', linewidth=0.5)

        for n, i in enumerate(reversed(bits)):
            line = [((x >> i) & 1) for x in data]
            plt.step(t, np.array(line) + 2*n, 'r', linewidth=2, where='post')
            plt.text(-0.5, 2*n, str(i))

        plt.xlim([t[0], t[-1]])
        plt.ylim([0, 2*len(bits)+1])

        plt.gca().axis('off')
        plt.show()

    while len(data) > 0:
        if len(data) > line_length:
            d = data[0:line_length]
            data = data[line_length:]
        else:
            d = data
            data = []

        _plot_timing_diagram(d, bits)

def plot_codeword_diagram(ts, cws, range=None):
    """
    Takes a list of timestamps (X) and codewords (Y) and produces a simple 'stem' plot of the two.
    The plot is useful for visually checking that the codewords are detected at regular intervals.
    Can also be used for visual verification of standard codeword patterns such as the staircase used
    for calibration.
    """
    plt.figure(figsize=(20, 10))
    plt.stem((np.array(ts)-ts[0])*10.0/3, np.array(cws))
    if range is not None:
        plt.xlim(range[0], range[1])
        xticks = np.arange(range[0], range[1], step=20)
        while len(xticks) > 20:
            xticks = xticks[::2]
        plt.xticks(xticks)
    plt.xlabel('Time (ns)')
    plt.ylabel('Codeword (#)')
    plt.grid()
    plt.show()

def _gen_set_cmd(dev_set_func, node_path: str):
    """
    Generates a set function based on the dev_set_type method (e.g., seti)
    and the node_path (e.g., '/dev8003/sigouts/1/mode'
    """
    def set_cmd(val):
        return dev_set_func(node_path, val)
    return set_cmd

def _gen_get_cmd(dev_get_func, node_path: str):
    """
    Generates a get function based on the dev_set_type method (e.g., geti)
    and the node_path (e.g., '/dev8003/sigouts/1/mode'
    """
    def get_cmd():
        return dev_get_func(node_path)
    return get_cmd

##########################################################################
# Exceptions
##########################################################################

class ziDAQError(Exception):
    """Exception raised when no DAQ has been connected."""
    pass

class ziModuleError(Exception):
    """Exception raised when a module generates an error."""
    pass

class ziValueError(Exception):
    """Exception raised when a wrong or empty value is returned."""
    pass

class ziCompilationError(Exception):
    """Exception raised when an AWG program fails to compile."""
    pass

class ziDeviceError(Exception):
    """Exception raised when a class is used with the wrong device type."""
    pass

class ziOptionsError(Exception):
    """Exception raised when a device does not have the right options installed."""
    pass

class ziVersionError(Exception):
    """Exception raised when a device does not have the right firmware versions."""
    pass

class ziReadyError(Exception):
    """Exception raised when a device was started which is not ready."""
    pass

class ziRuntimeError(Exception):
    """Exception raised when a device detects an error at runtime."""
    pass

class ziConfigurationError(Exception):
    """Exception raised when a wrong configuration is detected."""
    pass

##########################################################################
# Mock classes
##########################################################################

class MockDAQServer():
    """
    This class implements a mock version of the DAQ object used for
    communicating with the instruments. It contains dummy declarations of
    the most important methods implemented by the server and used by
    the instrument drivers.

    Important: The Mock server creates some default 'nodes' (basically
    just entries in a 'dict') based on the device name that is used when
    connecting to a device. These nodes differ depending on the instrument
    type, which is determined by the number in the device name: dev2XXX are
    UHFQA instruments, dev8XXX are HDAWG8 instruments, dev10XXX are PQSC
    instruments.
    """

    def __init__(self, server, port, apilevel, verbose=False):
        self.server = server
        self.port = port
        self.apilevel = apilevel
        self.device = None
        self.interface = None
        self.nodes = {'/zi/devices/connected': {'type': 'String', 'value': ''}}
        self.devtype = None
        self.poll_nodes = []
        self.verbose = verbose

    def awgModule(self):
        return MockAwgModule(self)

    def setDebugLevel(self, debuglevel: int):
        print('Setting debug level to {}'.format(debuglevel))

    def connectDevice(self, device, interface):
        if self.device is not None:
            raise ziDAQError(
                'Trying to connect to a device that is already connected!')

        if self.interface is not None and self.interface != interface:
            raise ziDAQError(
                'Trying to change interface on an already connected device!')

        self.device = device
        self.interface = interface

        if self.device.lower().startswith('dev2'):
            self.devtype = 'UHFQA'
        elif self.device.lower().startswith('dev8'):
            self.devtype = 'HDAWG8'
        elif self.device.lower().startswith('dev10'):
            self.devtype = 'PQSC'

        # Add paths
        filename = os.path.join(os.path.dirname(os.path.abspath(
            __file__)), 'zi_parameter_files', 'node_doc_{}.json'.format(self.devtype))
        if not os.path.isfile(filename):
            raise ziRuntimeError(
                'No parameter file available for devices of type ' + self.devtype)
        # NB: defined in parent class
        self._load_parameter_file(filename=filename)

        # Update connected status
        self.nodes['/zi/devices/connected']['value'] = self.device

        # Set the LabOne revision
        self.nodes['/zi/about/revision'] = {'type': 'Integer', 'value': 200802104}

        self.nodes[f'/{self.device}/features/devtype'] = {'type': 'String', 'value': self.devtype}
        self.nodes[f'/{self.device}/system/fwrevision'] = {'type': 'Integer', 'value': 99999}
        self.nodes[f'/{self.device}/system/fpgarevision'] = {'type': 'Integer', 'value': 99999}
        self.nodes[f'/{self.device}/system/slaverevision'] = {'type': 'Integer', 'value': 99999}

        if self.devtype == 'UHFQA':
            self.nodes[f'/{self.device}/features/options'] = {'type': 'String', 'value': 'QA\nAWG'}
            for i in range(16):
                self.nodes[f'/{self.device}/awgs/0/waveform/waves/{i}'] = {'type': 'ZIVectorData', 'value': np.array([])}
            for i in range(10):
                self.nodes[f'/{self.device}/qas/0/integration/weights/{i}/real'] = {'type': 'ZIVectorData', 'value': np.array([])}
                self.nodes[f'/{self.device}/qas/0/integration/weights/{i}/imag'] = {'type': 'ZIVectorData', 'value': np.array([])}
                self.nodes[f'/{self.device}/qas/0/result/data/{i}/wave'] = {'type': 'ZIVectorData', 'value': np.array([])}
            self.nodes[f'/{self.device}/raw/dios/0/delay'] = {'type': 'Integer', 'value': 0}
            self.nodes[f'/{self.device}/dios/0/extclk'] = {'type': 'Integer', 'value': 0}
            self.nodes[f'/{self.device}/dios/0/drive'] = {'type': 'Integer', 'value': 0}
            self.nodes[f'/{self.device}/dios/0/mode'] = {'type': 'Integer', 'value': 0}
        elif self.devtype == 'HDAWG8':
            self.nodes[f'/{self.device}/features/options'] = {'type': 'String', 'value': 'PC\nME'}
            self.nodes[f'/{self.device}/raw/error/json/errors'] = {
                'type': 'String', 'value': '{"sequence_nr" : 0, "new_errors" : 0, "first_timestamp" : 0, "timestamp" : 0, "timestamp_utc" : "2019-08-07 17 : 33 : 55", "messages" : []}'}
            for i in range(32):
                self.nodes['/' + self.device +
                        '/raw/dios/0/delays/' + str(i) + '/value'] = {'type': 'Integer', 'value': 0}
            self.nodes[f'/{self.device}/raw/error/blinkseverity'] = {'type': 'Integer', 'value': 0}
            self.nodes[f'/{self.device}/raw/error/blinkforever'] = {'type': 'Integer', 'value': 0}
            self.nodes[f'/{self.device}/dios/0/extclk'] = {'type': 'Integer', 'value': 0}
            for awg_nr in range(4):
                for i in range(2048):
                    self.nodes[f'/{self.device}/awgs/{awg_nr}/waveform/waves/{i}'] = {
                        'type': 'ZIVectorData', 'value': np.array([])}
                    self.nodes[f'/{self.device}/awgs/{awg_nr}/waveform/waves/{i}'] = {
                        'type': 'ZIVectorData', 'value': np.array([])}
                    self.nodes[f'/{self.device}/awgs/{awg_nr}/waveform/waves/{i}'] = {
                        'type': 'ZIVectorData', 'value': np.array([])}
                    self.nodes[f'/{self.device}/awgs/{awg_nr}/waveform/waves/{i}'] = {
                        'type': 'ZIVectorData', 'value': np.array([])}
            for sigout_nr in range(8):
                self.nodes[f'/{self.device}/sigouts/{sigout_nr}/precompensation/fir/coefficients'] = {
                    'type': 'ZIVectorData', 'value': np.array([])}
            self.nodes[f'/{self.device}/dios/0/mode'] = {'type': 'Integer', 'value': 0}
            self.nodes[f'/{self.device}/dios/0/extclk'] = {'type': 'Integer', 'value': 0}
            self.nodes[f'/{self.device}/dios/0/drive'] = {'type': 'Integer', 'value': 0}
            for dio_nr in range(32):
                self.nodes[f'/{self.device}/raw/dios/0/delays/{dio_nr}/value'] = {'type': 'Integer', 'value': 0}
            self.nodes[f'/{self.device}/raw/error/clear'] = {
                'type': 'Integer', 'value': 0}
        elif self.devtype == 'PQSC':
            self.nodes[f'/{self.device}/raw/error/json/errors'] = {
                'type': 'String', 'value': '{"sequence_nr" : 0, "new_errors" : 0, "first_timestamp" : 0, "timestamp" : 0, "timestamp_utc" : "2019-08-07 17 : 33 : 55", "messages" : []}'}
            self.nodes[f'/{self.device}/raw/error/clear'] = {
                'type': 'Integer', 'value': 0}

    def listNodesJSON(self, path):
        pass

    def getString(self, path):
        if path not in self.nodes:
            raise ziRuntimeError("Unknown node '" + path +
                                 "' used with mocked server and device!")

        if self.nodes[path]['type'] != 'String':
            raise ziRuntimeError(
                "Trying to node '" + path + "' as string, but the type is '" + self.nodes[path]['type'] + "'!")

        return self.nodes[path]['value']

    def getInt(self, path):
        if path not in self.nodes:
            raise ziRuntimeError("Unknown node '" + path +
                                 "' used with mocked server and device!")

        if self.verbose:
            print('getInt', path, int(self.nodes[path]['value']))

        return int(self.nodes[path]['value'])

    def getDouble(self, path):
        if path not in self.nodes:
            raise ziRuntimeError("Unknown node '" + path +
                                 "' used with mocked server and device!")

        if self.verbose:
            print('getDouble', path, float(self.nodes[path]['value']))

        return float(self.nodes[path]['value'])

    def setInt(self, path, value):
        if path not in self.nodes:
            raise ziRuntimeError("Unknown node '" + path +
                                 "' used with mocked server and device!")

        if self.verbose:
            print('setInt', path, value)

        self.nodes[path]['value'] = value

    def setDouble(self, path, value):
        if path not in self.nodes:
            raise ziRuntimeError("Unknown node '" + path +
                                 "' used with mocked server and device!")
        if self.verbose:
            print('setDouble', path, value)
        self.nodes[path]['value'] = value

    def setVector(self, path, value):
        if path not in self.nodes:
            raise ziRuntimeError("Unknown node '" + path +
                                 "' used with mocked server and device!")

        if self.nodes[path]['type'] != 'ZIVectorData':
            raise ziRuntimeError("Unable to set node '" + path + "' of type " +
                                 self.nodes[path]['type'] + " using setVector!")

        self.nodes[path]['value'] = value

    def setComplex(self, path, value):
        if path not in self.nodes:
            raise ziRuntimeError("Unknown node '" + path +
                                 "' used with mocked server and device!")

        if not self.nodes[path]['type'].startswith('Complex'):
            raise ziRuntimeError("Unable to set node '" + path + "' of type " +
                                 self.nodes[path]['type'] + " using setComplex!")

        if self.verbose:
            print('setComplex', path, value)
        self.nodes[path]['value'] = value

    def getComplex(self, path):
        if path not in self.nodes:
            raise ziRuntimeError("Unknown node '" + path +
                                 "' used with mocked server and device!")

        if not self.nodes[path]['type'].startswith('Complex'):
            raise ziRuntimeError("Unable to get node '" + path + "' of type " +
                                 self.nodes[path]['type'] + " using getComplex!")

        if self.verbose:
            print('getComplex', path, self.nodes[path]['value'])
        return self.nodes[path]['value']

    def get(self, path, flat, flags):
        if path not in self.nodes:
            raise ziRuntimeError("Unknown node '" + path +
                                 "' used with mocked server and device!")

        return {path: [{'vector': self.nodes[path]['value']}]}

    def getAsEvent(self, path):
        self.poll_nodes.append(path)

    def poll(self, poll_time, timeout, flags, flat):
        poll_data = {}

        for path in self.poll_nodes:
            if self.verbose:
                print('poll', path)
            m = re.match(r'/(\w+)/qas/0/result/data/(\d+)/wave', path)
            if m:
                poll_data[path] = [{'vector': np.random.rand(
                    self.getInt('/' + m.group(1) + '/qas/0/result/length'))}]
                continue

            m = re.match(r'/(\w+)/qas/0/monitor/inputs/(\d+)/wave', path)
            if m:
                poll_data[path] = [{'vector': np.random.rand(
                    self.getInt('/' + m.group(1) + '/qas/0/monitor/length'))}]
                continue

            m = re.match(r'/(\w+)/awgs/(\d+)/ready', path)
            if m:
                poll_data[path] = {'value': [1]}
                continue

            poll_data[path] = {'value': [0]}

        return poll_data

    def subscribe(self, path):
        if self.verbose:
            print('subscribe', path)

        self.poll_nodes.append(path)

    def unsubscribe(self, path):
        if self.verbose:
            print('unsubscribe', path)

        if path in self.poll_nodes:
            self.poll_nodes.remove(path)

    def sync(self):
        """The sync method does not need to do anything as there are no
        device delays to deal with when using the mock server.
        """
        pass

    def _load_parameter_file(self, filename: str):
        """
        Takes in a node_doc JSON file auto generates paths based on
        the contents of this file.
        """
        with open(filename) as fo:
            f = fo.read()
        node_pars = json.loads(f)
        for par in node_pars.values():
            node = par['Node'].split('/')
            # The parfile is valid for all devices of a certain type
            # so the device name has to be split out.
            parpath = '/' + self.device + '/' + '/'.join(node)
            if par['Type'].startswith('Integer'):
                self.nodes[parpath.lower()] = {'type': par['Type'], 'value': 0}
            elif par['Type'].startswith('Double'):
                self.nodes[parpath.lower()] = {
                    'type': par['Type'], 'value': 0.0}
            elif par['Type'].startswith('Complex'):
                self.nodes[parpath.lower()] = {
                    'type': par['Type'], 'value': 0 + 0j}
            elif par['Type'].startswith('String'):
                self.nodes[parpath.lower()] = {
                    'type': par['Type'], 'value': ''}

class MockAwgModule():
    """
    This class implements a mock version of the awgModule object used for
    compiling and uploading AWG programs. It doesn't actually compile anything, but
    only maintains a counter of how often the compilation method has been executed.

    For the future, the class could be updated to allow the user to select whether
    the next compilation should be successful or not in order to enable more
    flexibility in the unit tests of the actual drivers.
    """

    def __init__(self, daq):
        self._daq = daq
        self._device = None
        self._index = None
        self._sourcestring = None
        self._compilation_count = {}
        if not os.path.isdir('awg/waves'):
            os.makedirs('awg/waves')

    def get_compilation_count(self, index):
        if index not in self._compilation_count:
            raise ziModuleError(
                'Trying to access compilation count of invalid index ' + str(index) + '!')

        return self._compilation_count[index]

    def set(self, path, value):
        if path == 'awgModule/device':
            self._device = value
        elif path == 'awgModule/index':
            self._index = value
            if self._index not in self._compilation_count:
                self._compilation_count[self._index] = 0
        elif path == 'awgModule/compiler/sourcestring':
            # The compiled program is stored in _sourcestring
            self._sourcestring = value
            if self._index not in self._compilation_count:
                raise ziModuleError(
                    'Trying to compile AWG program, but no AWG index has been configured!')

            if self._device is None:
                raise ziModuleError(
                    'Trying to compile AWG program, but no AWG device has been configured!')

            self._compilation_count[self._index] += 1
            self._daq.setInt('/' + self._device + '/' +
                             'awgs/' + str(self._index) + '/ready', 1)

    def get(self, path):
        if path == 'awgModule/device':
            value = [self._device]
        elif path == 'awgModule/index':
            value[self._index]
        elif path == 'awgModule/compiler/statusstring':
            value = ['File successfully uploaded']
        else:
            value = ['']

        for elem in reversed(path.split('/')[1:]):
            rv = {elem: value}
            value = rv

        return rv

    def execute(self):
        pass

##########################################################################
# Class
##########################################################################

class ZI_base_instrument(Instrument):
    """
    This is a base class for Zurich Instruments instrument drivers.
    It includes functionality that is common to all instruments. It maintains
    a list of available nodes as JSON files in the 'zi_parameter_files'
    subfolder. The parameter files should be regenerated when newer versions
    of the firmware are installed on the instrument.

    The base class also manages waveforms for the instruments. The waveforms
    are kept in a table, which is kept synchronized with CSV files in the
    awg/waves folder belonging to LabOne. The base class will select whether
    to compile and configure an instrument based on changes to the waveforms
    and to the requested AWG program. Basically, if a waveform changes length
    or if the AWG program changes, then the program will be compiled and
    uploaded the next time the user executes the 'start' method. If a waveform
    has changed, but the length is the same, then the waveform will simply
    be updated on the instrument using a a fast waveform upload technique. Again,
    this is triggered when the 'start' method is called.
    """

    ##########################################################################
    # Constructor
    ##########################################################################

    def __init__(self,
                 name: str,
                 device: str,
                 interface: str= '1GbE',
                 server: str= 'localhost',
                 port: int= 8004,
                 apilevel: int= 5,
                 num_codewords: int= 0,
                 awg_module: bool=True,
                 logfile: str = None,
                 **kw) -> None:
        """
        Input arguments:
            name:           (str) name of the instrument as seen by the user
            device          (str) the name of the device e.g., "dev8008"
            interface       (str) the name of the interface to use ('1GbE' or 'USB')
            server          (str) the host where the ziDataServer is running
            port            (int) the port to connect to for the ziDataServer (don't change)
            apilevel        (int) the API version level to use (don't change unless you know what you're doing)
            awg_module      (bool) create an awgModule
            num_codewords   (int) the number of codeword-based waveforms to prepare
            logfile         (str) file name where all commands should be logged
        """
        t0 = time.time()
        super().__init__(name=name, **kw)

        # Decide which server to use based on name
        if server == 'emulator':
            log.info('Connecting to mock DAQ server')
            self.daq = MockDAQServer(server, port, apilevel)
        else:
            log.info('Connecting to DAQ server')
            self.daq = zi.ziDAQServer(server, port, apilevel)

        if not self.daq:
            raise(ziDAQError())

        self.daq.setDebugLevel(0)

        # Handle absolute path
        self.use_setVector = "setVector" in dir(self.daq)

        # Connect a device
        if not self._is_device_connected(device):
            log.info(f'Connecting to device {device}')
            self.daq.connectDevice(device, interface)
        self.devname = device
        self.devtype = self.gets('features/devtype')

        # We're now connected, so do some sanity checking
        self._check_devtype()
        self._check_versions()
        self._check_options()

        # Default waveform length used when initializing waveforms to zero
        self._default_waveform_length = 32

        # add qcodes parameters based on JSON parameter file
        # FIXME: we might want to skip/remove/(add  to _params_to_skip_update) entries like AWGS/*/ELF/DATA,
        #       AWGS/*/SEQUENCER/ASSEMBLY, AWGS/*/DIO/DATA
        filename = os.path.join(os.path.dirname(os.path.abspath(
            __file__)), 'zi_parameter_files', 'node_doc_{}.json'.format(self.devtype))
        if not os.path.isfile(filename):
            log.info(f"{self.devname}: Parameter file not found, creating  '{filename}''")
            self._create_parameter_file(filename=filename)

        try:
            # NB: defined in parent class
            log.info(f'{self.devname}: Loading parameter file')
            self._load_parameter_file(filename=filename)
        except FileNotFoundError:
            # Should never happen as we just created the file above
            log.error(f"{self.devname}: parameter file for data parameters {filename} not found")
            raise

        # Create modules
        if awg_module:
            self._awgModule = self.daq.awgModule()
            self._awgModule.set('awgModule/device', device)
            self._awgModule.execute()

            # Will hold information about all configured waveforms
            self._awg_waveforms = {}

            # Asserted when AWG needs to be reconfigured
            self._awg_needs_configuration = [False]*(self._num_channels()//2)
            self._awg_program = [None]*(self._num_channels()//2)

            # Create waveform parameters
            self._num_codewords = 0
            # CH: this Delft function should not be needed for us, and removing it
            # should save time in the init script. Please let me know if you
            # experience any issues with the init of HDAWG/UHF.
            # self._add_codeword_waveform_parameters(num_codewords)
        else:
            self._awgModule = None

        # Create other neat parameters
        self._add_extra_parameters()
        # A list of all subscribed paths
        self._subscribed_paths = []

        self._awg_source_strings = {}

        # Structure for storing errors
        self._errors = None
        # Structure for storing errors that should be demoted to warnings
        self._errors_to_ignore = []
        # Make initial error check
        self.check_errors()

        # Optionally setup log file
        if logfile is not None:
            self._logfile = open(logfile, 'w')
        else:
            self._logfile = None

        # Show some info
        serial = self.get('features_serial')
        options = self.get('features_options')
        fw_revision = self.get('system_fwrevision')
        fpga_revision = self.get('system_fpgarevision')
        log.info('{}: serial={}, options={}, fw_revision={}, fpga_revision={}'
                 .format(self.devname, serial, options.replace('\n', '|'), fw_revision, fpga_revision))

        self.connect_message(begin_time=t0)

    ##########################################################################
    # Private methods: Abstract Base Class methods
    ##########################################################################

    def _check_devtype(self):
        """
        Checks that the driver is used with the correct device-type.
        """
        raise NotImplementedError('Virtual method with no implementation!')

    def _check_options(self):
        """
        Checks that the correct options are installed on the instrument.
        """
        raise NotImplementedError('Virtual method with no implementation!')

    def _check_versions(self):
        """
        Checks that sufficient versions of the firmware are available.
        """
        raise NotImplementedError('Virtual method with no implementation!')

    def _check_awg_nr(self, awg_nr):
        """
        Checks that the given AWG index is valid for the device.
        """
        raise NotImplementedError('Virtual method with no implementation!')

    def _update_num_channels(self):
        raise NotImplementedError('Virtual method with no implementation!')

    def _update_awg_waveforms(self):
        raise NotImplementedError('Virtual method with no implementation!')

    def _num_channels(self):
        raise NotImplementedError('Virtual method with no implementation!')

    def _add_extra_parameters(self) -> None:
        """
        Adds extra useful parameters to the instrument.
        """
        log.info(f'{self.devname}: Adding extra parameters')
        self.add_parameter(
            'timeout',
            unit='s',
            initial_value=30,
            parameter_class=ManualParameter,
            vals=validators.Ints())

    ##########################################################################
    # Private methods
    ##########################################################################

    def _add_codeword_waveform_parameters(self, num_codewords) -> None:
        """
        Adds parameters that are used for uploading codewords.
        It also contains initial values for each codeword to ensure
        that the "upload_codeword_program" works.
        """
        docst = ('Specifies a waveform for a specific codeword. ' +
                 'The waveforms must be uploaded using ' +
                 '"upload_codeword_program". The channel number corresponds' +
                 ' to the channel as indicated on the device (1 is lowest).')

        self._params_to_skip_update = []
        log.info(f'{self.devname}: Adding codeword waveform parameters')
        for ch in range(self._num_channels()):
            for cw in range(max(num_codewords, self._num_codewords)):
                # NB: parameter naming identical to QWG
                wf_name = gen_waveform_name(ch, cw)

                if cw >= self._num_codewords and wf_name not in self.parameters:
                    # Add parameter
                    self.add_parameter(
                        wf_name,
                        label='Waveform channel {} codeword {:03}'.format(
                            ch+1, cw),
                        vals=validators.Arrays(),  # min_value, max_value = unknown
                        set_cmd=self._gen_write_waveform(ch, cw),
                        get_cmd=self._gen_read_waveform(ch, cw),
                        docstring=docst)
                    self._params_to_skip_update.append(wf_name)
                    # Make sure the waveform data is up-to-date
                    self._gen_read_waveform(ch, cw)()
                elif cw >= num_codewords:
                    # Delete parameter as it's no longer needed
                    if wf_name in self.parameters:
                        self.parameters.pop(wf_name)
                        self._awg_waveforms.pop(wf_name)

        # Update the number of codewords
        self._num_codewords = num_codewords

    def _load_parameter_file(self, filename: str):
        """
        Takes in a node_doc JSON file auto generates parameters based on
        the contents of this file.
        """
        with open(filename) as fo:
            f = fo.read()
        node_pars = json.loads(f)
        for par in node_pars.values():
            node = par['Node'].split('/')
            # The parfile is valid for all devices of a certain type
            # so the device name has to be split out.
            parname = '_'.join(node).lower()
            parpath = '/' + self.devname + '/' + '/'.join(node)

            # This block provides the mapping between the ZI node and QCoDes
            # parameter.
            par_kw = {}
            par_kw['name'] = parname
            if par['Unit'] != 'None':
                par_kw['unit'] = par['Unit']
            else:
                par_kw['unit'] = 'arb. unit'

            par_kw['docstring'] = par['Description']
            if "Options" in par.keys():
                # options can be done better, this is not sorted
                par_kw['docstring'] += '\nOptions:\n' + str(par['Options'])

            # Creates type dependent get/set methods
            if par['Type'] == 'Integer (64 bit)':
                par_kw['set_cmd'] = _gen_set_cmd(self.seti, parpath)
                par_kw['get_cmd'] = _gen_get_cmd(self.geti, parpath)
                # min/max not implemented yet for ZI auto docstrings #352
                par_kw['vals'] = validators.Ints()

            elif par['Type'] == 'Integer (enumerated)':
                par_kw['set_cmd'] = _gen_set_cmd(self.seti, parpath)
                par_kw['get_cmd'] = _gen_get_cmd(self.geti, parpath)
                par_kw['vals'] = validators.Ints(min_value=0,
                                                 max_value=len(par["Options"]))

            elif par['Type'] == 'Double':
                par_kw['set_cmd'] = _gen_set_cmd(self.setd, parpath)
                par_kw['get_cmd'] = _gen_get_cmd(self.getd, parpath)
                # min/max not implemented yet for ZI auto docstrings #352
                par_kw['vals'] = validators.Numbers()

            elif par['Type'] == 'Complex Double':
                par_kw['set_cmd'] = _gen_set_cmd(self.setc, parpath)
                par_kw['get_cmd'] = _gen_get_cmd(self.getc, parpath)
                # min/max not implemented yet for ZI auto docstrings #352
                par_kw['vals'] = validators.Anything()

            elif par['Type'] == 'ZIVectorData':
                par_kw['set_cmd'] = _gen_set_cmd(self.setv, parpath)
                par_kw['get_cmd'] = _gen_get_cmd(self.getv, parpath)
                # min/max not implemented yet for ZI auto docstrings #352
                par_kw['vals'] = validators.Arrays()

            elif par['Type'] == 'String':
                par_kw['set_cmd'] = _gen_set_cmd(self.sets, parpath)
                par_kw['get_cmd'] = _gen_get_cmd(self.gets, parpath)
                par_kw['vals'] = validators.Strings()

            elif par['Type'] == 'CoreString':
                par_kw['get_cmd'] = _gen_get_cmd(self.getd, parpath)
                par_kw['set_cmd'] = None  # Not implemented
                par_kw['vals'] = validators.Strings()

            elif par['Type'] == 'ZICntSample':
                par_kw['get_cmd'] = None  # Not implemented
                par_kw['set_cmd'] = None  # Not implemented
                par_kw['vals'] = None  # Not implemented

            elif par['Type'] == 'ZITriggerSample':
                par_kw['get_cmd'] = None  # Not implemented
                par_kw['set_cmd'] = None  # Not implemented
                par_kw['vals'] = None  # Not implemented

            elif par['Type'] == 'ZIDIOSample':
                par_kw['get_cmd'] = None  # Not implemented
                par_kw['set_cmd'] = None  # Not implemented
                par_kw['vals'] = None  # Not implemented

            elif par['Type'] == 'ZIAuxInSample':
                par_kw['get_cmd'] = None  # Not implemented
                par_kw['set_cmd'] = None  # Not implemented
                par_kw['vals'] = None  # Not implemented

            elif par['Type'] == 'ZIScopeWave':
                par_kw['get_cmd'] = None  # Not implemented
                par_kw['set_cmd'] = None  # Not implemented
                par_kw['vals'] = None  # Not implemented

            else:
                raise NotImplementedError(
                    "Parameter '{}' of type '{}' not supported".format(
                        parname, par['Type']))

            # If not readable/writable the methods are removed after the type
            # dependent loop to keep this more readable.
            if 'Read' not in par['Properties']:
                par_kw['get_cmd'] = None
            if 'Write' not in par['Properties']:
                par_kw['set_cmd'] = None
            self.add_parameter(**par_kw)

    def _create_parameter_file(self, filename: str):
        """
        This generates a json file Containing the node_docs as extracted
        from the ZI instrument API.

        Replaces the use of the s_node_pars and d_node_pars files.
        """
        # Get all interesting nodes
        nodes = json.loads(self.daq.listNodesJSON('/' + self.devname))

        modified_nodes = {}

        # Do some name mangling
        for name, node in nodes.items():
            name = name.replace('/' + self.devname.upper() + '/', '')
            node['Node'] = name
            modified_nodes[name] = node

        # Dump the nodes
        with open(filename, "w") as json_file:
            json.dump(modified_nodes, json_file, indent=4, sort_keys=True)

    def _is_device_connected(self, device):
        """
        Return true if the given device is already connected to the server.
        """
        if device.lower() in [x.lower() for x in self.daq.getString('/zi/devices/connected').split(',')]:
            return True
        else:
            return False

    def _get_full_path(self, paths):
        """
        Concatenates the device name with one or more paths to create a fully
        qualified path for use in the server.
        """
        if type(paths) is list:
            for p, n in enumerate(paths):
                if p[0] != '/':
                    paths[n] = ('/' + self.devname + '/' + p).lower()
                else:
                    paths[n] = paths[n].lower()
        else:
            if paths[0] != '/':
                paths = ('/' + self.devname + '/' + paths).lower()
            else:
                paths = paths.lower()

        return paths

    def _get_awg_directory(self):
        """
        Returns the AWG directory where waveforms should be stored.
        """
        return os.path.join(self._awgModule.get('awgModule/directory')['directory'][0], 'awg')

    def _initialize_waveform_to_zeros(self):
        """
        Generates all zeros waveforms for all codewords.
        """
        t0 = time.time()
        wf = np.zeros(self._default_waveform_length)
        waveform_params = [value for key, value in self.parameters.items()
                           if 'wave_ch' in key.lower()]
        for par in waveform_params:
            par(wf)
        t1 = time.time()
        log.debug(
            'Set all waveforms to zeros in {:.1f} ms'.format(1.0e3*(t1-t0)))

    def _gen_write_waveform(self, ch, cw):
        def write_func(waveform):
            log.debug(f"{self.devname}: Writing waveform (len {len(waveform)}) to ch{ch} cw{cw}")
            # Determine which AWG this waveform belongs to
            awg_nr = ch//2

            # Name of this waveform
            wf_name = gen_waveform_name(ch, cw)

            # Check that we're allowed to modify this waveform
            if self._awg_waveforms[wf_name]['readonly']:
                raise ziConfigurationError(
                    'Trying to modify read-only waveform on '
                    'codeword {}, channel {}'.format(cw, ch))

            # The length of HDAWG waveforms should be a multiple of 8 samples.
            if (len(waveform) % 8) != 0:
                log.debug(f"{self.devname}: waveform is not a multiple of 8 samples, appending zeros.")
                extra_zeros = 8-(len(waveform) % 8)
                waveform = np.concatenate([waveform, np.zeros(extra_zeros)])

            # If the length has changed, we need to recompile the AWG program
            if len(waveform) != len(self._awg_waveforms[wf_name]['waveform']):
                log.debug(f"{self.devname}: Length of waveform has changed. Flagging awg as requiring recompilation.")
                self._awg_needs_configuration[awg_nr] = True

            # Update the associated CSV file
            log.debug(f"{self.devname}: Updating csv waveform {wf_name}, for ch{ch}, cw{cw}")
            self._write_csv_waveform(ch=ch, cw=cw, wf_name=wf_name,
                                     waveform=waveform)

            # And the entry in our table and mark it for update
            self._awg_waveforms[wf_name]['waveform'] = waveform
            log.debug(f"{self.devname}: Marking waveform as dirty.")
            self._awg_waveforms[wf_name]['dirty'] = True

        return write_func

    def _write_csv_waveform(self, ch: int, cw: int, wf_name: str, waveform) -> None:
        filename = os.path.join(
            self._get_awg_directory(), 'waves',
            self.devname + '_' + wf_name + '.csv')
        np.savetxt(filename, waveform, delimiter=",")

    def _gen_read_waveform(self, ch, cw):
        def read_func():
            # AWG
            awg_nr = ch//2

            # Name of this waveform
            wf_name = gen_waveform_name(ch, cw)
            log.debug(f"{self.devname}: Reading waveform {wf_name} for ch{ch} cw{cw}")
            # Check if the waveform data is in our dictionary
            if wf_name not in self._awg_waveforms:
                log.debug(f"{self.devname}: Waveform not in self._awg_waveforms: reading from csv file.")
                # Initialize elements
                self._awg_waveforms[wf_name] = {
                    'waveform': None, 'dirty': False, 'readonly': False}
                # Make sure everything gets recompiled
                log.debug(f"{self.devname}: Flagging awg as requiring recompilation.")
                self._awg_needs_configuration[awg_nr] = True
                # It isn't, so try to read the data from CSV
                waveform = self._read_csv_waveform(ch, cw, wf_name)
                # Check whether  we got something
                if waveform is None:
                    log.debug(f"{self.devname}: Waveform CSV does not exist, initializing to zeros.")
                    # Nope, initialize to zeros
                    waveform = np.zeros(32)
                    self._awg_waveforms[wf_name]['waveform'] = waveform
                    # write the CSV file
                    self._write_csv_waveform(ch, cw, wf_name, waveform)
                else:
                    # Got data, update dictionary
                    self._awg_waveforms[wf_name]['waveform'] = waveform

            # Get the waveform data from our dictionary, which must now
            # have the data
            return self._awg_waveforms[wf_name]['waveform']

        return read_func

    def _read_csv_waveform(self, ch: int, cw: int, wf_name: str):
        filename = os.path.join(
            self._get_awg_directory(), 'waves',
            self.devname + '_' + wf_name + '.csv')
        try:
            log.debug(f"{self.devname}: reading waveform from csv '{filename}'")
            return np.genfromtxt(filename, delimiter=',')
        except OSError as e:
            # if the waveform does not exist yet dont raise exception
            log.warning(e)
            return None

    def _length_match_waveforms(self, awg_nr):
        """
        Adjust the length of a codeword waveform such that each individual
        waveform of the pair has the same length
        """
        log.info('Length matching waveforms for dynamic waveform upload.')
        wf_table = self._get_waveform_table(awg_nr)

        matching_updated = False
        iter_id = 0
        # We iterate over the waveform table
        while(matching_updated or iter_id == 0):

            iter_id += 1
            if iter_id > 10:
                raise StopIteration
            log.info('Length matching iteration {}.'.format(iter_id))
            matching_updated = False

            for wf_name, other_wf_name in wf_table:
                len_wf = len(self._awg_waveforms[wf_name]['waveform'])
                len_other_wf = len(self._awg_waveforms[other_wf_name]['waveform'])

                # First one is shorter
                if len_wf < len_other_wf:
                    log.info(f"{self.devname}: Modifying {wf_name} for length matching.")
                    # Temporarily unset the readonly flag to be allowed to append zeros
                    readonly = self._awg_waveforms[wf_name]['readonly']
                    self._awg_waveforms[wf_name]['readonly'] = False
                    self.set(wf_name, np.concatenate(
                        (self._awg_waveforms[wf_name]['waveform'], np.zeros(len_other_wf-len_wf))))
                    self._awg_waveforms[wf_name]['dirty'] = True
                    self._awg_waveforms[wf_name]['readonly'] = readonly
                    matching_updated = True
                elif len_other_wf < len_wf:
                    log.info(f"{self.devname}: Modifying {other_wf_name} for length matching.")
                    readonly = self._awg_waveforms[other_wf_name]['readonly']
                    self._awg_waveforms[other_wf_name]['readonly'] = False
                    self.set(other_wf_name, np.concatenate(
                        (self._awg_waveforms[other_wf_name]['waveform'], np.zeros(len_wf-len_other_wf))))
                    self._awg_waveforms[other_wf_name]['dirty'] = True
                    self._awg_waveforms[other_wf_name]['readonly'] = readonly
                    matching_updated = True

    def _clear_dirty_waveforms(self, awg_nr):
        """
        Adjust the length of a codeword waveform such that each individual
        waveform of the pair has the same length
        """
        log.info(f"{self.devname}: Clearing dirty waveform tag for AWG {awg_nr}")
        for cw in range(self._num_codewords):
            wf_name = gen_waveform_name(2*awg_nr+0, cw)
            self._awg_waveforms[wf_name]['dirty'] = False

            other_wf_name = gen_waveform_name(2*awg_nr+1, cw)
            self._awg_waveforms[other_wf_name]['dirty'] = False

    def _clear_readonly_waveforms(self, awg_nr):
        """
        Clear the read-only flag of all configured waveforms. Typically used when switching
        configurations (i.e. programs).
        """
        for cw in range(self._num_codewords):
            wf_name = gen_waveform_name(2*awg_nr+0, cw)
            self._awg_waveforms[wf_name]['readonly'] = False

            other_wf_name = gen_waveform_name(2*awg_nr+1, cw)
            self._awg_waveforms[other_wf_name]['readonly'] = False

    def _set_readonly_waveform(self, ch: int, cw: int):
        """
        Mark a waveform as being read-only. Typically used to limit which waveforms the user
        is allowed to change based on the overall configuration of the instrument and the type
        of AWG program being executed.
        """
        # Sanity check
        if cw >= self._num_codewords:
            raise ziConfigurationError(
                'Codeword {} is out of range of the configured number of codewords ({})!'.format(cw, self._num_codewords))

        if ch >= self._num_channels():
            raise ziConfigurationError(
                'Channel {} is out of range of the configured number of channels ({})!'.format(ch, self._num_channels()))

        # Name of this waveform
        wf_name = gen_waveform_name(ch, cw)

        # Check if the waveform data is in our dictionary
        if wf_name not in self._awg_waveforms:
            raise ziConfigurationError(
                'Trying to mark waveform {} as read-only, but the waveform has not been configured yet!'.format(wf_name))

        self._awg_waveforms[wf_name]['readonly'] = True

    def _upload_updated_waveforms(self, awg_nr):
        """
        Loop through all configured waveforms and use dynamic waveform uploading
        to update changed waveforms on the instrument as needed.
        """
        # Fixme. the _get_waveform_table should also be implemented for the UFH
        log.info(f"{self.devname}: Using dynamic waveform update for AWG {awg_nr}.")
        wf_table = self._get_waveform_table(awg_nr)

        for dio_cw, (wf_name, other_wf_name) in enumerate(wf_table):
            if self._awg_waveforms[wf_name]['dirty'] or self._awg_waveforms[other_wf_name]['dirty']:
                # Combine the waveforms and upload
                wf_data = merge_waveforms(self._awg_waveforms[wf_name]['waveform'],
                                          self._awg_waveforms[other_wf_name]['waveform'])
                # Write the new waveform
                self.setv(
                    'awgs/{}/waveform/waves/{}'.format(awg_nr, dio_cw), wf_data)

    def _codeword_table_preamble(self, awg_nr):
        """
        Defines a snippet of code to use in the beginning of an AWG program in order to define the waveforms.
        The generated code depends on the instrument type. For the HDAWG instruments, we use the setDIOWaveform
        function. For the UHF-QA we simply define the raw waveforms.
        """
        raise NotImplementedError('Virtual method with no implementation!')

    def _configure_awg_from_variable(self, awg_nr):
        """
        Configures an AWG with the program stored in the object in the self._awg_program[awg_nr] member.
        """
        log.info(f"{self.devname}: Configuring AWG {awg_nr} with predefined codeword program")
        if self._awg_program[awg_nr] is not None:
            full_program = \
                '// Start of automatically generated codeword table\n' + \
                self._codeword_table_preamble(awg_nr) + \
                '// End of automatically generated codeword table\n' + \
                self._awg_program[awg_nr]

            self.configure_awg_from_string(awg_nr, full_program)
        else:
            logging.info(f"{self.devname}: No program configured for awg_nr {awg_nr}.")

    def _write_cmd_to_logfile(self, cmd):
        if self._logfile is not None:
            now = datetime.now()
            now_str = now.strftime("%d/%m/%Y %H:%M:%S")
            self._logfile.write(f'#{now_str}\n')
            self._logfile.write(f'{self.name}.{cmd}\n')

    def _flush_logfile(self):
        if self._logfile is not None:
            self._logfile.flush()

    ##########################################################################
    # Public methods: node helpers
    ##########################################################################

    def setd(self, path, value) -> None:
        self._write_cmd_to_logfile(f'daq.setDouble("{path}", {value})')
        self.daq.setDouble(self._get_full_path(path), value)

    def getd(self, path):
        return self.daq.getDouble(self._get_full_path(path))

    def seti(self, path, value) -> None:
        self._write_cmd_to_logfile(f'daq.setDouble("{path}", {value})')
        self.daq.setInt(self._get_full_path(path), value)

    def geti(self, path):
        return self.daq.getInt(self._get_full_path(path))

    def sets(self, path, value) -> None:
        self._write_cmd_to_logfile(f'daq.setString("{path}", {value})')
        self.daq.setString(self._get_full_path(path), value)

    def gets(self, path):
        return self.daq.getString(self._get_full_path(path))

    def setc(self, path, value) -> None:
        self._write_cmd_to_logfile(f'daq.setComplex("{path}", {value})')
        self.daq.setComplex(self._get_full_path(path), value)

    def getc(self, path):
        return self.daq.getComplex(self._get_full_path(path))

    def setv(self, path, value) -> None:
        # Handle absolute path
        if self.use_setVector:
            self._write_cmd_to_logfile(f'daq.setVector("{path}", np.array({np.array2string(value, separator=",")}))')
            self.daq.setVector(self._get_full_path(path), value)
        else:
            self._write_cmd_to_logfile(f'daq.vectorWrite("{path}", np.array({np.array2string(value, separator=",")}))')
            self.daq.vectorWrite(self._get_full_path(path), value)

    def getv(self, path):
        path = self._get_full_path(path)
        value = self.daq.get(path, True, 0)
        if path not in value:
            raise ziValueError('No value returned for path ' + path)
        else:
            return value[path][0]['vector']

    def getdeep(self, path, timeout=5.0):
        path = self._get_full_path(path)

        self.daq.getAsEvent(path)
        while timeout > 0.0:
            value = self.daq.poll(0.01, 500, 4, True)
            if path in value:
                return value[path]
            else:
                timeout -= 0.01

        return None

    def subs(self, path:str) -> None:
        full_path = self._get_full_path(path)
        if full_path not in self._subscribed_paths:
            self._subscribed_paths.append(full_path)
        self.daq.subscribe(full_path)

    def unsubs(self, path:str=None) -> None:
        if path is None:
            for path in self._subscribed_paths:
                self.daq.unsubscribe(path)
            self._subscribed_paths.clear()
        else:
            full_path = self._get_full_path(path)
            if full_path in self._subscribed_paths:
                del self._subscribed_paths[self._subscribed_paths.index(full_path)]
            self.daq.unsubscribe(full_path)

    def poll(self, poll_time=0.1):
        # The timeout of 1ms (second argument) is smaller than in the Delft
        # driver version (500ms) to allow fast spectroscopy.
        return self.daq.poll(poll_time, 1, 4, True)

    def sync(self) -> None:
        self.daq.sync()

    ##########################################################################
    # Public methods
    ##########################################################################

    def start(self, **kw):
        """ Start the sequencer
        :param kw: currently ignored, added for compatibilty with other
            instruments that accept kwargs in start().
        """
        log.info(f"{self.devname}: Starting '{self.name}'")
        self.check_errors()

        # Loop through each AWG and check whether to reconfigure it
        for awg_nr in range(self._num_channels()//2):
            self._length_match_waveforms(awg_nr)

            # If the reconfiguration flag is set, upload new program
            if self._awg_needs_configuration[awg_nr]:
                log.debug(f"{self.devname}: Detected awg configuration tag for AWG {awg_nr}.")
                self._configure_awg_from_variable(awg_nr)
                self._awg_needs_configuration[awg_nr] = False
                self._clear_dirty_waveforms(awg_nr)
            else:
                log.debug(f"{self.devname}: Did not detect awg configuration tag for AWG {awg_nr}.")
                # Loop through all waveforms and update accordingly
                self._upload_updated_waveforms(awg_nr)
                self._clear_dirty_waveforms(awg_nr)

        # Start all AWG's
        for awg_nr in range(self._num_channels()//2):
            # Skip AWG's without programs
            if self._awg_program[awg_nr] is None:
                # to configure all awgs use "upload_codeword_program" or specify
                # another program
                logging.info(f"{self.devname}: Not starting awg_nr {awg_nr}.")
                continue
            # Check that the AWG is ready
            if not self.get('awgs_{}_ready'.format(awg_nr)):
                raise ziReadyError(
                    'Tried to start AWG {} that is not ready!'.format(awg_nr))
            # Enable it
            self.set('awgs_{}_enable'.format(awg_nr), 1)
        log.info(f"{self.devname}: Started '{self.name}'")

    def stop(self):
        log.info('Stopping {}'.format(self.name))
        # Stop all AWG's
        for awg_nr in range(self._num_channels()//2):
            self.set('awgs_{}_enable'.format(awg_nr), 0)

        self.check_errors()

    # FIXME: temporary solution for issue
    def FIXMEclose(self) -> None:
        try:
            # Disconnect application server
            self.daq.disconnect()
        except AttributeError:
            pass
        super().close()

    def check_errors(self, errors_to_ignore=None) -> None:
        raise NotImplementedError('Virtual method with no implementation!')

    def clear_errors(self) -> None:
        raise NotImplementedError('Virtual method with no implementation!')

    def demote_error(self, code: str):
        """
        Demote a ZIRuntime error to a warning.

        Arguments
            code (str)
            The error code of the exception to ignore.
            The error code gets logged as an error before the exception
            is raised. The code is a string like "DIOCWCASE".
        """
        self._errors_to_ignore.append(code)

    def reset_waveforms_zeros(self):
        """
        Sets all waveforms to an array of 48 zeros.
        """
        t0 = time.time()

        wf = np.zeros(48)
        waveform_params = [value for key, value in self.parameters.items()
                           if 'wave_ch' in key.lower()]
        for par in waveform_params:
            par(wf)

        t1 = time.time()
        log.info('Set all waveforms to zeros in {:.1f} ms'.format(1.0e3*(t1-t0)))

    def configure_awg_from_string(self, awg_nr: int, program_string: str,
                                  timeout: float=15):
        """
        Uploads a program string to one of the AWGs in a UHF-QA or AWG-8.

        This function is tested to work and give the correct error messages
        when compilation fails.
        """
        log.info(f'{self.devname}: Configuring AWG {awg_nr} from string.')
        # Check that awg_nr is set in accordance with devtype
        self._check_awg_nr(awg_nr)

        self._awg_source_strings[awg_nr] = program_string

        t0 = time.time()
        success_and_ready = False
        if not hasattr(self, 'compiler_statusstring'):
            self.compiler_statusstring = ''

        # This check (and while loop) is added as a workaround for #9
        while not success_and_ready:
            new_statusstring = f'{self.devname}: Configuring AWG {awg_nr}...'
            log.info(new_statusstring)
            self.compiler_statusstring += new_statusstring

            self._awgModule.set('awgModule/index', awg_nr)
            self._write_cmd_to_logfile(f"_awgModule.set('awgModule/index', {awg_nr})")
            self._awgModule.set(
                'awgModule/compiler/sourcestring', program_string)
            self._write_cmd_to_logfile(f"_awgModule.set('awgModule/compiler/sourcestring', \'\'\'{program_string}\'\'\')")

            succes_msg = 'File successfully uploaded'

            # Success is set to False when either a timeout or a bad compilation
            # message is encountered.
            success = True
            while len(self._awgModule.get('awgModule/compiler/sourcestring')
                      ['compiler']['sourcestring'][0]) > 0:
                time.sleep(0.01)

                if (time.time()-t0 >= timeout):
                    success = False
                    raise TimeoutError(
                        'Timeout while waiting for compilation to finish!')

            comp_msg = (self._awgModule.get(
                'awgModule/compiler/statusstring')['compiler']
                ['statusstring'][0])

            if not comp_msg.endswith(succes_msg):
                success = False

            if not success:
                print("Compilation failed, printing program:")
                for i, line in enumerate(program_string.splitlines()):
                    print(i+1, '\t', line)
                print('\n')
                raise ziCompilationError(comp_msg)

            # Give the device one second to respond
            for i in range(10):
                ready = self.getdeep(
                    'awgs/{}/ready'.format(awg_nr))['value'][0]
                if ready != 1:
                    log.warning('AWG {} not ready'.format(awg_nr))
                    time.sleep(1)
                else:
                    success_and_ready = True
                    break

        t1 = time.time()
        new_statusstring = (self._awgModule.get(
            'awgModule/compiler/statusstring')
              ['compiler']['statusstring'][0] + ' in {:.2f}s'.format(t1-t0))
        log.info(new_statusstring)
        self.compiler_statusstring += new_statusstring

        # Check status
        if self.get('awgs_{}_waveform_memoryusage'.format(awg_nr)) > 1.0:
            log.warning(f'{self.devname}: Waveform memory usage exceeds available internal memory!')

        if self.get('awgs_{}_sequencer_memoryusage'.format(awg_nr)) > 1.0:
            log.warning(f'{self.devname}: Sequencer memory usage exceeds available instruction memory!')

    def plot_dio_snapshot(self, bits=range(32)):
        raise NotImplementedError('Virtual method with no implementation!')

    def plot_awg_codewords(self, awg_nr=0, range=None):
        raise NotImplementedError('Virtual method with no implementation!')

    def get_idn(self) -> dict:
        idn_dict = {}
        idn_dict['vendor'] = 'ZurichInstruments'
        idn_dict['model'] = self.devtype
        idn_dict['serial'] = self.devname
        idn_dict['firmware'] = self.geti('system/fwrevision')
        idn_dict['fpga_firmware'] = self.geti('system/fpgarevision')

        return idn_dict

    def load_default_settings(self):
        raise NotImplementedError('Virtual method with no implementation!')

    def assure_ext_clock(self) -> None:
        raise NotImplementedError('Virtual method with no implementation!')


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

import math
import random
import sys
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.path as path
import numpy as np
import itertools

def generate_histo(array, name):
	plt.hist(array, bins=1000)
	plt.title('Sampling Importance Resampling')
	plt.show()

def ZGJ(regn, fx, init, itmx, step_function):

	ndim = len(regn[0])/2

	total_interval = [0 for i in range(ndim)]

	for di in range(ndim):
		pp = [[] for i in range(ndim)]
		pp_diff = [[] for i in range(ndim)]

		total_interval[di] = [[regn[0][di], regn[0][ndim + di]]]
		area = 0

		for it in range(itmx):
			xo = regn[0][di] # xo = 0.0
			xn = regn[0][ndim + di] # xn = 2.0
			tl_interval = []
			i = 0

			subarea = 0
			num_points = 100

			for interval in total_interval[di]: # 对每一个interval撒点，排序，记录权重，
				num_bin = 3
				i += 1

				num_points = (int(100*(interval[1] - interval[0])/(regn[0][ndim + di] - regn[0][di]))+1)
				points = interval[0] + (interval[1] - interval[0]) * np.random.rand(num_points) #随机点的个数一定要足够大（大于num_bin），否则可能导致点不够分，最后使得区间数目达不到要求
				points.sort(axis = 0)

				wgts = abs(fx(points))
				sum_wgts = np.sum(wgts)
				cut = sum_wgts/num_bin
				# print '当前区间的cut：', cut, i

				k = 0
				sum_wgt = 0
				xo = interval[0]

				for point in points:
					k += 1
					sum_wgt += abs(fx([point])) # 先考虑fx大于零的情况
					if sum_wgt > cut:
						tl_interval.append([xo, point])
						xo = point
						sum_wgt = 0
						k = 0

				tl_interval.append([xo, interval[1]])

			total_interval[di] = tl_interval

			print 'iteration: ', it + 1, len(total_interval)

	fxpp, subgrid_area = step_function(total_interval, fx, ndim)
	expectation = np.sum(np.array(fxpp) * np.array(subgrid_area))

	return expectation

def fx(x):
	# return x**2+2
	# return np.exp((x-0.5)**2)
	return np.cos(x[0])+1
	# return 2

def step_function(inter, fx, ndim):
	do = inter[0]
	for di in range(ndim - 1):
		do = itertools.product(do, inter[di + 1])

	grids = do
	xx = [0 for i in range(ndim)]

	pp = []
	subgrid_area = []

	for grid in grids:
		area_grid = 1
		for di in range(ndim):
			xx[di] = grid[di][0] + (grid[di][1] - grid[di][0]) * np.random.rand(10)
			area_grid *= (grid[di][1] - grid[di][0])

		subgrid_area.append(area_grid)

		xxx = [0 for n in range(10)]

		for n in range(10):
			xxx[n] = [xx[di][n] for di in range(ndim)]

		pp.append(np.mean([fx(xxx[n]) for n in range(10)])) # 每一格点里的函数平均值已经得到，同时被收集到pp里面

	return pp, subgrid_area

################### To Begin ###################

if __name__ == "__main__":
	init = 1
	itmx = 2
	EE = ZGJ([[0.0, 0.0, 2.0, 1.0]], fx, init, itmx, step_function) # regen should be float
	print 'iteration: ', itmx, 'times'
	print 'integral: ', EE

################### To-Do List ###################
# 1. damp function
# 2. 处理一下函数值可能为负数的情况
# 3. the more iteration times, the lower speed
# 4. sigma estimate
# 5. alpha running
# 6. multi-dimension