import time
from collections import deque
import cupy as cp
import copy
import multiprocessing
import cupyx.scipy.ndimage
import cupyx.scipy.fftpack as scipyfftpack
import numpy as np
import prysm
from scipy import interpolate, ndimage, fftpack, optimize, signal
import matplotlib.pyplot as plt

from lentil import wavefront_config__old as conf
from lentil.wavefront_config__old import SPACIAL_FREQS, BASE_WAVELENGTH, MODEL_WVLS, DEFAULT_SAMPLES
from lentil.constants_utils import *

zcache = prysm.zernike.zcache.regular
cupyzcache = {}
numpyzcache_arrays = {}
cupyzcache_arrays = {}
cache_idx = {}
cache_idx["cp"] = None
cache_idx["np"] = None
window_cache = {}

all_but_z4_and_z9_phasecache = dict(np={}, cp={})

RETURN_MTF = 1
RETURN_OTF = 2
RETURN_LSF = 3
RETURN_PSF = 4
RETURN_WITH_PROCESSING_DETAILS = 5

settings_cache = {}

mask_cache = deque(maxlen=conf.MASK_CACHE_SIZE)

deltas = np.linspace(-1,1, 19) * 1e-12
a = np.random.random((256, 256))
psf_size = 256
mses = []
for d in deltas:
    normshift_x = d
    normshift_y = 0
    zoom_factor = 1.0000
    transform = np.array(((1.0 / zoom_factor, 0), (0, 1.0 / zoom_factor)))
    offset_x = (psf_size - 1 + 1) / 2 * (1.0 - 1.0 / zoom_factor - normshift_x)
    offset_y = (psf_size - 1 + 1) / 2 * (1.0 - 1.0 / zoom_factor - normshift_y)
    order = conf.PSF_SPLINE_ORDER
    zoomed_mono_psf = ndimage.affine_transform(a, transform, (offset_y, offset_x),
                                       order=order)
    mses.append(((zoomed_mono_psf - a)**2).sum())
    print(d, mses[-1])

# plt.plot(deltas, mses)
# plt.show()
# exit()
class TestSettings:
    def __init__(self, defocus, p):
        self.defocus = defocus
        self.p = p
        self.mono = False
        self.plot = False
        self.dummy = False
        self.allow_cuda = conf.USE_CUDA
        self.id_or_hash = 0
        self.strehl_estimate = 1.0
        self.fftsize = None
        self.phasesamples = None
        self.return_otf = True
        self.return_otf_mtf = False
        self.return_psf = False
        self.return_prysm_mtf = False
        self.return_mask = False
        self.mask = None
        self.prysm_mtf = None
        self.cpu_gpu_arraysize_boundary = conf.CPU_GPU_ARRAYSIZE_BOUNDARY
        self.effective_q = None
        self.guide_mtf = None
        self.q_autosize_scalar = conf.Q_AUTOSIZE_SCALAR
        self.phase_autosize_scalar = conf.PHASE_AUTOSIZE_SCALAR
        self.cache_sizes = True
        self.x_loc = IMAGE_WIDTH / 2
        self.y_loc = IMAGE_HEIGHT / 2
        self.pixel_vignetting = True
        self.lens_vignetting = True
        self.default_exit_pupil_position_mm = 100
        self.exif = None
        self.fix_pupil_rotation = True

    @property
    def fftshape(self):
        if self.fftsize is None:
            return None
        return self.fftsize, self.fftsize

    @property
    def phaseshape(self):
        if self.phasesamples is None:
            return None
        return self.phasesamples, self.phasesamples

    @property
    def is_valid(self):
        if self.fftsize is None:
            return False
        if self.phasesamples is None:
            return False
        if self.p is None:
            return False
        if self.defocus is None:
            return False
        return True

    def get_processing_details(self):
        get_processing_details(self)
        return self

class TestResults:
    def __init__(self):
        self.lsf = None
        self.psf = None
        self.otf = None
        self.timings = None
        self.strehl = None
        self.fftsize = None
        self.samples = None
        self.id_or_hash = None
        self.used_cuda = None

    def copy_important_settings(self, s: TestSettings):
        self.fftsize = s.fftsize
        self.samples = s.phasesamples
        self.id_or_hash = s.id_or_hash

    def get_mtf(self):
        return abs(self.otf[0]), abs(self.otf[1])

def zoom(inarr, xfactor=1.0, yfactor=1.0, xoffset=0.0, yoffset=0.0, affine_transform=ndimage.affine_transform, me=np):
    shape = inarr.shape
    transform = me.array(((1.0 / yfactor, 0), (0, 1.0 / xfactor)))
    offset_x = (shape[0]) / 2 * (1.0 - 1.0 / xfactor) - xoffset / xfactor
    offset_y = (shape[1]) / 2 * (1.0 - 1.0 / yfactor) - yoffset / yfactor
    if me is cp:
        return affine_transform(inarr, transform, (offset_y, offset_x),
                                           order=1)
    else:
        real = affine_transform(inarr.real, transform, (offset_y, offset_x),
                                           order=conf.PSF_SPLINE_ORDER)
        imag = affine_transform(inarr.imag, transform, (offset_y, offset_x),
                                           order=conf.PSF_SPLINE_ORDER)
        return real + 1j * imag

def get_z9(p, modelwavelength):
    rel_wv = modelwavelength / BASE_WAVELENGTH
    spca = p.get('spca', 0.0) * 30
    spca2 = p.get('spca2', 0.0) * 30

    spcaz9 = (modelwavelength / BASE_WAVELENGTH - 1.0) * spca + spca * 0.028
    spca2z9 = (rel_wv - 1.0) ** 2 * spca2 * 10 - spca2 * 0.06
    return (p.get('z9', 0.0) + spcaz9 + spca2z9) * conf.BASE_WAVELENGTH

def get_z4(defocus, p, modelwavelength):
    fstop_base_ratio = p['fstop'] / p['base_fstop']
    rel_wv = modelwavelength / BASE_WAVELENGTH
    loca = p.get('loca', 0.0) * 30
    loca1 = p.get('loca1', 0.0) * 30

    locadefocus = (rel_wv - 1.0) ** 2 * 10 * loca - loca * 0.06
    loca1defocus = (rel_wv - 1.0) * 1 * loca1 + loca1 * 0.027
    base_z4 = ((defocus - p.get('df_offset', 0)) * p.get('df_step', 1)) * fstop_base_ratio ** 2
    # print(p)
    return -(base_z4 - locadefocus - loca1defocus) * conf.BASE_WAVELENGTH

def get_lca_shifts(s: TestSettings, modelwavelength, samplespacing):
    rel_wv = modelwavelength / 0.54

    img_height = calc_image_height(s.x_loc, s.y_loc)

    px = s.p.get('tca_slr', 0.0) * 1e2 * img_height
    py = 0

    shiftx = (rel_wv - 1.0) ** 2 * px * 10 - px / 14
    shifty = (rel_wv - 1.0) ** 2 * py * 10 - py / 14
    return shiftx / samplespacing / s.fftsize, shifty / samplespacing / s.fftsize

strehl_estimate_cache = None
mtf_cache = None
return_cache = None

def get_used_zernikes(iterable, cache=True, me=np):
    max = 0
    arr = me.zeros(48, dtype="int")
    idx = me.zeros(48, dtype="int") - 1
    used = []
    count = 0
    for item in iterable:
        if item[0].lower() == "z" and item[1].isdigit():
            print(item)
            zn = int(item[1:])
            arr[zn - 1] = 1
            if (zn) > max:
                max = zn
            idx[zn - 1] = count
            count += 1
            used.append(zn)
    arr[3] = 1
    print(arr, max, idx, 99)
    return arr, max, idx, used

def get_processing_details(s: TestSettings):
    if s.id_or_hash is not None and s.id_or_hash in settings_cache:
        stup = settings_cache[s.id_or_hash]
        if s.fftsize is None:
            s.fftsize = stup[0]
        if s.phasesamples is None:
            s.phasesamples = stup[1]
        if s.effective_q is None:
            s.effective_q = stup[2]
        s.allow_cuda = s.allow_cuda and s.fftsize > s.cpu_gpu_arraysize_boundary
        return s

    # if s.return_type == RETURN_LSF:
    #     s.fftsize = 128
    #     s.phasesamples = 64
    #     s.effective_q = 2
    #     return s

    minimum_q = np.clip((s.strehl_estimate * 4) * s.q_autosize_scalar, 2, 3)
    # minimum_q = np.clip((0.5 + s.strehl_estimate * 3.5) * s.q_autosize_scalar, 1.0, 5)
    min_samples = -np.inf
    f_stopped_down = s.p['fstop'] / s.p['base_fstop']
    if s.guide_mtf is None:
        min_samples = 384
    else:
        for otf in s.guide_mtf:
            freqs = np.arange(0, 65) / 64
            zero_plus_spacial_freqs = np.concatenate(([0], SPACIAL_FREQS, [1.0, 2.0]))
            interpotf_real = interpolate.InterpolatedUnivariateSpline(zero_plus_spacial_freqs, np.concatenate(([1.0], otf.real, [0,0])), k=2)(freqs)
            interpotf_imag = interpolate.InterpolatedUnivariateSpline(zero_plus_spacial_freqs, np.concatenate(([1.0], otf.imag, [0,0])), k=2)(freqs)
            interpotf = interpotf_real + 1j * interpotf_imag
            fftin = np.concatenate((interpotf[:-1], np.flip(interpotf[1:])))
            lsfshifted = np.abs(fftpack.ifft(fftin))
            lsfmax = np.maximum(lsfshifted[:64], np.flip(lsfshifted[64:]))
            lsfmax /= lsfmax.max()

            fitweights = np.clip((0.1 - lsfmax)*25, 0, 0.999) ** 4
            fitweights = (0.01 < lsfmax) * (lsfmax < 0.12)
            x_arr = np.arange(len(lsfmax))

            def cost(params, return_curve=False):
                a, b = params
                c = 0
                expcurve = b * np.exp(-0.1 * a * x_arr) + c
                if return_curve:
                    return expcurve
                return ((lsfmax - expcurve) ** 2 * fitweights).mean()

            a, b = optimize.minimize(cost, (1.0, 1.0,), bounds=((0.01, 70), (0.1, 30),)).x
            c = 0
            cutoff = 0.03

            needed_width = -10 / a * np.log((cutoff - c) / b)

            # plt.plot(cost((a, b), return_curve=True), label="fit")
            # plt.plot(lsfshifted / lsfshifted.max(), label="lsfshifted")
            # plt.plot(fftin, label="fftin")
            # plt.plot(lsfmax, label="lsfmax")
            # plt.ylim(0, 1)
            # plt.hlines([cutoff], 0, 64)
            # plt.legend()
            # plt.show()

            min_samples_this_axis = needed_width * s.phase_autosize_scalar * 9
            if min_samples_this_axis > min_samples:
                min_samples = min_samples_this_axis
    # min_samples = 64 + (1.0 - minmtf[3])**2 * 400

    # samples = int(min_samples / 2 + 1) * 2

    # for samples in CUDA_GOOD_FFT_SIZES: # limit sizes for cache reasons
    #     if samples >= min_samples:
    #         break
    for power in range(4, 10):
        samples = 2 ** power
        if samples > min_samples:
            break
        samples = int((2 ** power * 1.5) / 2 + 0.5) * 2
        if samples > min_samples:
            break

    effective_q_without_padding = f_stopped_down

    min_fftsize = minimum_q * samples / effective_q_without_padding

    for fftsize in (CUDA_GOOD_FFT_SIZES if s.allow_cuda else CPU_GOOD_FFT_SIZES):
        if fftsize >= min_fftsize:
            break

    # s.allow_cuda = s.allow_cuda and (samples + fftsize) >= (s.cpu_gpu_arraysize_boundary / 2)
    s.allow_cuda = s.allow_cuda and fftsize >= s.cpu_gpu_arraysize_boundary

    effective_q = fftsize / samples * effective_q_without_padding

    assert (fftsize - samples) % 2 == 0
    assert fftsize % 2 == 0
    assert samples % 2 == 0

    if s.fftsize is None:
        s.fftsize = fftsize
    if s.phasesamples is None:
        s.phasesamples = samples

    s.effective_q = effective_q

    if s.id_or_hash is not None:
        settings_cache[s.id_or_hash] = fftsize, samples, effective_q

    # s.fftsize = 1024
    # s.phasesamples = 512

    return s

# s = TestSettings(0, dict(fstop=5.6, base_fstop=1.4))
# s.phasesamples = 256
# get_processing_details(s)
# exit()

def try_wavefront(s: TestSettings):
    global window_cache

    orig_s = copy.copy(s)
    orig_s.p = s.p.copy()
    t = time.time()

    tr = TestResults()
    tr.copy_important_settings(s)
    if s.id_or_hash == 1:
        print("step", s.p['df_step'])

    if s.dummy:
        return tr

    if s.p['fstop'] < s.p['base_fstop']:
        raise ValueError("Base_fstop must be wider (lower) than fstop {} < {}".format(s.p['fstop'], s.p['base_fstop']))
    mul = 1
    mtfs = []
    pupilslices = []
    bestpupil = (np.inf, None, None)

    if not s.is_valid:
        s = get_processing_details(s)

    use_cuda = s.allow_cuda
    # prysm_path_q = 4
    # s.fftsize = prysm_path_q * s.phasesamples
    prysm_path_q = max(1, s.fftsize / s.phasesamples)

    if use_cuda:
        fft2 = cupyx.scipy.fftpack.fft2
        fftpack = cupyx.scipy.fftpack
        affine_transform = cupyx.scipy.ndimage.affine_transform
        # affine_transform = ndimage.affine_transform
    else:
        fft2 = scipyfftpack.fft2
        fftpack = scipyfftpack
        affine_transform = ndimage.affine_transform

    tr.used_cuda = use_cuda

    SAM_RADIOMETRIC_MODEL = True

    USE_PSF = False

    # cudadevice = cp.cuda.Device(0)

    def sync():
        pass
        # if use_cuda:
        #     cudadevice.synchronize()

    realdtype = "float32" if conf.PRECISION == 32 else "float64"
    complexdtype = "complex64" if conf.PRECISION == 32 else "complex128"

    me, engine_string = (cp, 'np') if use_cuda else (np, 'np')

    eval_wavelengths = [BASE_WAVELENGTH] if s.mono else MODEL_WVLS
    # eval_wavelengths = [BASE_WAVELENGTH] if s.mono else [0.5, 0.9]

    # s.p['loca'] *= 0.1
    # s.p['loca1'] *= 0.1
    # s.p['spca2'] *= 0.1
    # s.p['spca'] *= 0.1
    # s.p['tca_slr'] = 0
    if s.return_psf or USE_PSF or s.return_prysm_mtf:
        mono_psf_stack = me.zeros(s.fftshape, dtype="float64")

    psf_sag = np.zeros((s.fftsize, ), dtype="float64")
    psf_tan = np.zeros((s.fftsize, ), dtype="float64")
    psf_lst = []
    samplelst = []

    polychromatic_weights = np.array([float(photopic_fn(wv * 1e3) * d50_interpolator(wv)) for wv in eval_wavelengths])
    # polychromatic_weights **= 2
    # polychromatic_weights = me.array([0.5, 0.5])

    t_misc = 0
    t_pupils = 0
    t_get_phases = 0
    t_get_fcns = 0
    t_fcntransforms = 0
    t_pads = 0
    t_ffts = 0
    t_cudasyncs = 0
    t_affines = 0
    t_mtfs = 0
    t_init = time.time() - t
    t = time.time()

    mask = mask_pupil(s, engine=me, dtype=realdtype)

    if s.return_mask:
        tr.mask = mask

    t_maskmaking = time.time() - t

    used_zernikes_flags, max_zernike, cache_idx_, used_zernikes = get_used_zernikes(s.p.keys())

    if cache_idx[engine_string] is not None:
        if not me.all(cache_idx[engine_string] == cache_idx_):
            raise Exception("Zernicke cache does not match P dict")
    else:
        cache_idx[engine_string] = cache_idx_

    zkwargs = {}
    for key, value in s.p.items():

        # if key.upper() == 'Z9':
        #     zkwargs[key] = z9
        #     continue
        # if key.upper() == 'Z4':
        #     raise ValueError("No Z4 separately!")
        if key.lower().startswith('z') and key[1].isdigit():
            # if key not in ['z9', 'z16', 'z25', 'z36']:
            #     continue
            zkwargs[key] = value * mul * conf.BASE_WAVELENGTH

    z_arr_no_z4_z9 = np.zeros(used_zernikes_flags.sum(), dtype=realdtype)
    for key, value in zkwargs.items():
        idx = cache_idx['np'][int(key[1:]) - 1]
        z_arr_no_z4_z9[idx] = value
    z_arr_no_z4_z9[cache_idx[engine_string][3]] = 0
    z_arr_no_z4_z9[cache_idx[engine_string][8]] = 0
    print(zkwargs)
    print(zkwargs)
    print(zkwargs)
    print(zkwargs)
    print(z_arr_no_z4_z9)
    print(z_arr_no_z4_z9)
    print(z_arr_no_z4_z9)
    print(z_arr_no_z4_z9)
    print(z_arr_no_z4_z9)
    print(z_arr_no_z4_z9)
    print(z_arr_no_z4_z9)
    zhash = hash(tuple(z_arr_no_z4_z9))

    if me is not np:
        z_arr_no_z4_z9 = me.array(z_arr_no_z4_z9, dtype=realdtype)

    if me is cp:
        # cache = cupyzcache
        cache_array = cupyzcache_arrays
    else:
        # cache = zcache
        cache_array = numpyzcache_arrays

    for wvl_num, (model_wvl, polych_weight) in enumerate(zip(eval_wavelengths, polychromatic_weights)):
        t = time.time()
        rel_wv = model_wvl / BASE_WAVELENGTH

        z4 = get_z4(s.defocus, s.p, model_wvl)
        z9 = get_z9(s.p, model_wvl)

        samplelst.append(s.phasesamples)
        t_misc += time.time() - t
        if s.phasesamples not in cache_array:
            t = time.time()
            if s.phasesamples not in zcache:
                pupil = prysm.FringeZernike(used_zernikes_flags,
                                            dia=10, norm=False,
                                            wavelength=model_wvl,
                                            opd_unit="um",
                                            mask_target='none',
                                            samples=s.phasesamples, )  # This is just here to fill the coeff cache
            t_pupils += time.time() - t
            t = time.time()

            # if me is cp:
            #     cache[s.phasesamples] = {}
            cache_array[s.phasesamples] = me.empty((s.phasesamples, s.phasesamples, used_zernikes_flags.sum()), dtype=realdtype)
            # cache_array[s.phasesamples] = 1

            # for key, val in cache[s.phasesamples].items():
            if 1:
                idx = cache_idx['np'][key]
                # if me is cp:
                #     cache[s.phasesamples][key] = me.array(val)
                cache_array[s.phasesamples][:, :, idx] = me.array(val, dtype=realdtype)

                cache_idx[key] = idx
            print("graaagl")

            sync()
            t_get_phases += time.time() - t
        t = time.time()
        pupil = prysm.FringeZernike(dia=10, wavelength=model_wvl, norm=False,
                                    opd_unit="um",
                                    mask_target='none',
                                    samples=s.phasesamples, )
        t_pupils += time.time() - t
        t = time.time()

        if s.phasesamples not in all_but_z4_and_z9_phasecache[engine_string] or \
                all_but_z4_and_z9_phasecache[engine_string][s.phasesamples][0] != zhash:
            all_but_z4_and_z9_phasecache[engine_string][s.phasesamples] = zhash, cache_array[s.phasesamples] @ z_arr_no_z4_z9

        if me is cp:
            phase = cache_array[s.phasesamples][:, :, 3] * z4
            phase += all_but_z4_and_z9_phasecache[engine_string][s.phasesamples][1]
            phase += cache_array[s.phasesamples][:, :, 8] * z9
        else:
            # print(cache)
            phase = cache[s.phasesamples][3] * z4
            phase += all_but_z4_and_z9_phasecache[engine_string][s.phasesamples][1]
            phase += cache[s.phasesamples][8] * z9
        phase /= model_wvl
        sync()
        t_get_phases += time.time() - t

        t = time.time()
        # phase = me.array(pupil.change_phase_unit(to='waves', inplace=False),
        #                  dtype=realdtype)
        # phase = me.zeros((samples, samples), dtype=realdtype)
        wavefunction = me.exp(1j * 2 * me.pi * phase)
        wavefunction *= mask

        # if me is cp:
        #     plt.imshow(cp.asnumpy(me.angle(wavefunction)))
        #     plt.show()
        # else:
        #     plt.imshow(me.angle(wavefunction))
        #     plt.show()
        sync()
        t_get_fcns += time.time() - t

        if model_wvl == min(eval_wavelengths):
            if s.plot:
                mono_psf = prysm.PSF.from_pupil(pupil, efl=s.p['base_fstop'] * 10, Q=prysm_path_q)

            # psf_x_units, psf_y_units = prysm.propagation.prop_pupil_plane_to_psf_plane_units(wavefunction,
            #                                                              pupil.sample_spacing,
            #                                                              s.p['base_fstop'] * 10, model_wvl,
            #                                                              prysm_path_q)

            psf_sample_spacing = prysm.propagation.pupil_sample_to_psf_sample(pupil_sample=pupil.sample_spacing,
                                                          samples=s.fftsize,
                                                          wavelength=model_wvl,
                                                          efl=s.p['base_fstop'] * 10) * 1e-3
            psf_units = np.arange(-s.fftsize / 2, s.fftsize / 2) * psf_sample_spacing

        padpx = int((s.fftsize - s.phasesamples) / 2)

        ellip = s.p.get('ellip', 0)
        if ellip == 0:
            ellip = None
        else:
            xellip = np.clip(1 + ellip, 0.5, 1.0)
            yellip = np.clip(1.0 - ellip, 0.5, 1.0)

        if padpx > 0:
            pt = padpx, padpx
            if conf.ENABLE_PUPIL_DISTORTION and ellip is not None:
                t = time.time()
                wavefunction = zoom(wavefunction, xellip, yellip, affine_transform=affine_transform, me=me)
                t_fcntransforms += time.time() - t
            t = time.time()
            padded_cropped_pupil_fcn = me.pad(me.array(wavefunction, dtype=complexdtype), (pt, pt), mode="constant")
            t_pads += time.time() - t
        elif padpx < 0:
            t = time.time()
            padded_cropped_pupil_fcn = me.array(wavefunction[-padpx:-padpx+s.fftsize, -padpx:-padpx+s.fftsize], dtype=complexdtype)
            t_pads += time.time() - t
            if conf.ENABLE_PUPIL_DISTORTION and ellip is not None:
                t = time.time()
                padded_cropped_pupil_fcn = zoom(padded_cropped_pupil_fcn, xellip, yellip, affine_transform=affine_transform, me=me)
                t_fcntransforms += time.time() - t

        else:
            # padded_cropped_pupil_fcn = me.array(wavefunction, dtype=complexdtype)
            padded_cropped_pupil_fcn = wavefunction
            if conf.ENABLE_PUPIL_DISTORTION and ellip is not None:
                t = time.time()
                padded_cropped_pupil_fcn = zoom(padded_cropped_pupil_fcn, xellip, yellip, affine_transform=affine_transform, me=me)
                t_fcntransforms += time.time() - t
        t = time.time()
        pad = me.fft.fftshift(padded_cropped_pupil_fcn)
        sync()
        t_pads += time.time() - t
        t = time.time()

        # pad = me.fft.fft2(pad, norm='ortho')
        pad = fft2(pad, overwrite_x=True)
        shifted = me.fft.ifftshift(pad)
        impulse_response = me.absolute(shifted)
        impulse_response **= 2
        if not SAM_RADIOMETRIC_MODEL:
            impulse_response /= impulse_response.sum()
        impx = impulse_response.sum(axis=1)
        impy = impulse_response.sum(axis=0)

        sync()
        t_ffts += time.time() - t

        psf_size = impulse_response.shape[0]

        zoom_factor = float(np.clip(model_wvl / min(eval_wavelengths), 1.00001, np.inf))
        normshift_x, normshift_y = get_lca_shifts(s, model_wvl, psf_sample_spacing)

        if s.return_psf or USE_PSF:
            t = time.time()
            transform = me.array(((1.0 / zoom_factor, 0), (0, 1.0 / zoom_factor)))
            offset_x = (psf_size - 1 + 1) / 2 * (1.0 - 1.0 / zoom_factor - normshift_x)
            offset_y = (psf_size - 1 + 1) / 2 * (1.0 - 1.0 / zoom_factor - normshift_y)
            if use_cuda:
                order = 1
            else:
                order = conf.PSF_SPLINE_ORDER
            zoomed_mono_psf = affine_transform(impulse_response, transform, (offset_y, offset_x),
                                               order=order)
            if SAM_RADIOMETRIC_MODEL:
                zoomed_mono_psf *= polych_weight / zoomed_mono_psf.sum()
            else:
                zoomed_mono_psf *= polych_weight
            mono_psf_stack += zoomed_mono_psf
            sync()
            t_affines += time.time() - t
        # else:
        if 1:
            t = time.time()
            if me is cp:
                impx = cp.asnumpy(impx)
                impy = cp.asnumpy(impy)
            sync()
            t_cudasyncs += time.time() - t

            t = time.time()
            transform = np.array((1.0 / zoom_factor,))

            offset_x = (psf_size - 1 + 1) / 2 * (1.0 - 1.0 / zoom_factor - normshift_x)
            offset_y = (psf_size - 1 + 1) / 2 * (1.0 - 1.0 / zoom_factor - normshift_y)

            # impx /= impx.sum()
            # impy /= impy.sum()

            zoomx = ndimage.affine_transform(impx, transform, offset_x, order=conf.PSF_SPLINE_ORDER)
            zoomy = ndimage.affine_transform(impy, transform, offset_y, order=conf.PSF_SPLINE_ORDER)

            # zoomx = zoomed_mono_psf.mean(axis=0)
            # zoomy = zoomed_mono_psf.mean(axis=1)

            # plt.plot(zoomx / zoomx.max())
            # plt.plot(zoomx_ / zoomx_.max())
            # plt.show()

            # plt.plot(zoomy / zoomy.max())
            # plt.plot(zoomy_ / zoomy_.max())
            # plt.show()

            # print(model_wvl, s.id_or_hash)
            # plt.plot(psf_units, zoomx)
            # plt.plot(psf_units, zoomy)
            # plt.plot(psf_units, impx)
            # plt.plot(psf_units, impy)
            # plt.show()

            mul = 1.0 * zoom_factor
            # mul = 1

            if SAM_RADIOMETRIC_MODEL:
                # pass
                psf_sag += zoomx / zoomx.sum() * polych_weight
                psf_tan += zoomy / zoomy.sum() * polych_weight
            else:
                psf_sag += zoomx * polych_weight * mul
                psf_tan += zoomy * polych_weight * mul
            sync()
            t_affines += time.time() - t

        if s.plot:
            psf_lst.append(mono_psf)

        pupilslices.append(pupil.slice_x[1])

        metric = np.abs(rel_wv - 1)
        if metric < bestpupil[0]:
            bestpupil = metric, pupil, model_wvl, zkwargs

    t = time.time()

    if USE_PSF:
        psf_sag = mono_psf_stack.sum(axis=1)
        psf_tan = mono_psf_stack.sum(axis=0)
        if me is cp:
            psf_sag = cp.asnumpy(psf_sag)
            psf_tan = cp.asnumpy(psf_tan)
        # plt.plot(psf_sag / psf_sag.mean())
        # plt.plot(psf_sag_ / psf_sag_.mean())
        # plt.plot(((psf_sag / psf_sag.mean()) / (psf_sag_ / psf_sag_.mean()))[192:256+64])
        # plt.show()
    # ref_tr = _try_wavefront_prysmref(orig_s)

    if s.return_psf or s.return_prysm_mtf:
        npunits = cp.asnumpy(psf_units)
        prysm_psf = prysm.PSF(x=npunits, y=npunits, data=cp.asnumpy(mono_psf_stack))
        tr.psf = prysm_psf
        if 0:

            psf3 = prysm.PSF(x=npunits, y=npunits, data=cp.asnumpy(mono_psf_stack))
            psf2 = ref_tr.psf

            prysm_psf.data /= prysm_psf.data.sum()
            psf2.data /= psf2.data.sum()

            # assert np.allclose(prysm_psf.data, psf2.data)
            # assert np.allclose([prysm_psf.sample_spacing], [psf2.sample_spacing])
            psfnorm = prysm_psf.data / prysm_psf.data.sum()
            psf2norm = psf2.data / psf2.data.sum()
            psf3.data = np.clip(psfnorm / psf2norm, 0.3, 3)
            # print(psf3.data)
            # psf3.data[0,0] = max(psfnorm.max(), psf2norm.max())
            f, (a1, a2, a3) = plt.subplots(1, 3)

            tr.psf = prysm_psf
            prysm_psf.plot2d(axlim=18, ax=a1, fig=f)
            psf2.plot2d(axlim=18, ax=a2, fig=f)
            psf3.plot2d(axlim=18, ax=a3, fig=f)
            plt.show()
        # tr.otf = np.ones_like(SPACIAL_FREQS), np.ones_like(SPACIAL_FREQS)
        # return tr

    # plt.plot(psf_units, psf_sag / psf_sag.max()+ 0.1)
    # plt.plot(psf_units, mono_psf_stack.mean(axis=0) / mono_psf_stack.mean(axis=0).max())
    # plt.show()
    # plt.plot(psf_units, psf_tan / psf_tan.max() + 0.1)
    # plt.plot(psf_units, mono_psf_stack.mean(axis=1) / mono_psf_stack.mean(axis=1).max())
    # plt.show()

    # if s.return_type == RETURN_LSF:
    #     lsf_sag = mono_psf_stack.sum(axis=0)
    #     lsf_sag /= lsf_sag.max()
    #     lsf_tan = mono_psf_stack.sum(axis=1)
    #     lsf_tan /= lsf_tan.max()
    #     tr.lsf = lsf_sag, lsf_tan
    #     return tr

    centre = s.fftsize // 2
    mtf_x_units = prysm.fttools.forward_ft_unit(psf_sample_spacing * 1e-3, s.fftsize)
    sag_x = mtf_x_units[centre:]
    tan_x = mtf_x_units[centre:]

    if 0:
        if otf:
            mtf = me.fft.fftshift(me.fft.fft2(me.fft.ifftshift(mixedpsf)))
        # else:
    # mtf = me.absolute(me.fft.fft2(mono_psf_stack))
    # psf_y_units = prysm.fttools.forward_ft_unit((psf_x_units[1] - psf_x_units[0]) / 1e3, len(psf_y_units))
    # mtf = mtf / np.abs(mtf[0, 0])
    # sag_mod = mtf[0, :centre]
    # tan_mod = mtf[:centre, 0]

    # sag_mod = me.fft.fftshift(me.fft.fft(me.fft.ifftshift(psf_sag)))[centre:]
    # tan_mod = me.fft.fftshift(me.fft.fft(me.fft.ifftshift(psf_tan)))[centre:]

    # sag_mod /= np.abs(sag_mod)[0]
    # tan_mod /= np.abs(tan_mod)[0]

    mtf_mapper_fft_halfwindowsize_um = 16 * DEFAULT_PIXEL_SIZE * 1e6

    tukeykey = (s.fftsize, psf_sample_spacing)
    try:
        tukey_window = window_cache[tukeykey]
    except KeyError:
        if len(window_cache) > 300:
            window_cache = {}
        tukey_window = tukey(psf_units / mtf_mapper_fft_halfwindowsize_um, 0.6)
        window_cache[tukeykey] = tukey_window
        # print("PSF size {}um".format(psf_units[-1]*2), len(window_cache))

    # tukey_window = np.ones_like(psf_sag)

    # tukey_window = get_window(s.fftsize, psf_sample_spacing)
    # print("{:.3f} {:.3f}".format(np.sum(psf_sag * psf_units) / psf_sag.sum(), np.sum(psf_tan * psf_units) / psf_tan.sum()))
    sag_mod = normalised_centreing_fft(psf_sag * tukey_window, fftpack=scipyfftpack, engine=np)[:centre]
    tan_mod = normalised_centreing_fft(psf_tan * tukey_window, fftpack=scipyfftpack, engine=np)[:centre]

    # sag_mod = abs(fftpack.fft(np.fft.fftshift(psf_sag))[:centre])
    # tan_mod = abs(fftpack.fft(np.fft.fftshift(psf_tan))[:centre])
    # try:
    #     sag_mod /= sag_mod[0]
    # except FloatingPointError:
    #     sag_mod = np.zeros(centre)
    # try:
    #     tan_mod /= tan_mod[0]
    # except FloatingPointError:
    #     tan_mod = np.zeros(centre)
    # tan_mod /= tan_mod[0]

    # ref_mtf = ref_tr._prysm_mtf
    # (p_sag_x, p_sag_mod), (p_tan_x, p_tan_mod) = ref_mtf

    # assert np.allclose(p_sag_mod, abs(sag_mod))
    # assert np.allclose(p_tan_mod, abs(tan_mod))

    # plt.plot(p_sag_x, p_sag_mod)
    # plt.plot(sag_x, abs(sag_mod))
    # plt.show()
    # plt.plot(sag_x, abs(sag_mod) / p_sag_mod)
    # plt.show()
    # plt.plot(p_tan_x, p_tan_mod)
    # plt.plot(tan_x, abs(tan_mod))
    # plt.show()
    # plt.plot(tan_x, abs(tan_mod) / p_tan_mod)
    # plt.show()
    if s.return_prysm_mtf:
        prysm_mtf = prysm.MTF.from_psf(prysm_psf)
        tr.prysm_mtf = prysm_mtf

    if s.return_otf and s.return_otf_mtf:
        interpfn = interpolate.InterpolatedUnivariateSpline(sag_x, np.abs(sag_mod), k=1)
        sagmtf = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)
        interpfn = interpolate.InterpolatedUnivariateSpline(tan_x, np.abs(tan_mod), k=1)
        tanmtf = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)
        tr.otf = sagmtf, tanmtf

        # plt.plot(abs(tr.otf[0]))
        # plt.plot(ref_tr.otf[0])
        # plt.show()
        # plt.plot(abs(tr.otf[1]))
        # plt.plot(ref_tr.otf[1])
        # plt.show()

    if s.return_otf and not s.return_otf_mtf:
        interpfn = interpolate.InterpolatedUnivariateSpline(sag_x, np.real(sag_mod), k=1)
        sagmtf = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)
        interpfn = interpolate.InterpolatedUnivariateSpline(tan_x, np.real(tan_mod), k=1)
        tanmtf = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)
        interpfn = interpolate.InterpolatedUnivariateSpline(sag_x, np.imag(sag_mod), k=1)
        sagmtf_i = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)
        interpfn = interpolate.InterpolatedUnivariateSpline(tan_x, np.imag(tan_mod), k=1)
        tanmtf_i = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)
        # sagmtf = sagmtf + sagmtf_i * 1j
        # tanmtf = tanmtf + tanmtf_i * 1j

        # Normalise phase somehow

        # tr.otf = normalised_fft(sagmtf, sagmtf_i, SPACIAL_FREQS, inc_neg_freqs=False, return_type=COMPLEX_CARTESIAN), \
        #          normalised_fft(tanmtf, tanmtf_i, SPACIAL_FREQS, inc_neg_freqs=False, return_type=COMPLEX_CARTESIAN)

        # if s.id_or_hash == 0:
        #     normalised_fft(sagmtf, sagmtf_i, SPACIAL_FREQS, inc_neg_freqs=False,
        #                    return_type=COMPLEX_CARTESIAN, plot=True)
            # plt.plot(tuple(psf_sag))
            # plt.plot(tuple(psf_tan))
            # plt.show()
            # plt.plot(sagmtf, color='red')
            # plt.plot(tr.otf[0].real, '--', color='red')
            # plt.plot(tanmtf, color='green')
            # plt.plot(tr.otf[1].real, '--', color='green')
            # plt.plot(sagmtf_i, color='orange')
            # plt.plot(tr.otf[0].imag, '--', color='orange')
            # plt.plot(tanmtf_i, color='blue')
            # plt.plot(tr.otf[1].imag, '--', color='blue')
            # plt.show()

        tr.otf = sagmtf + 1j * sagmtf_i, tanmtf + 1j * tanmtf_i

    sync()
    t_mtfs += time.time() - t

    timings = dict(t_init=t_init,
                   t_maskmaking=t_maskmaking,
                   t_pupils=t_pupils,
                   t_get_phases=t_get_phases,
                   t_get_fcns=t_get_fcns,
                   t_fcntransforms=t_fcntransforms,
                   t_pads=t_pads,
                   t_ffts=t_ffts,
                   t_cudasyncs=t_cudasyncs,
                   t_affines=t_affines,
                   t_mtfs=t_mtfs,
                   t_misc=t_misc)

    tr.timings = timings

    sliceavg = me.average(me.array(pupilslices, dtype='float64'), axis=0, weights=polychromatic_weights)

    slice_ = bestpupil[1].slice_x[1]
    slicedv = me.abs(me.diff(sliceavg[me.isfinite(sliceavg)]))
    if slicedv.sum() != 0:
        peakiness = slicedv.max() / slicedv.mean()
    else:
        peakiness = 999.0
    strehl = float(bestpupil[1].strehl)

    if s.plot:
        # Plot each Z phase separately
        for key, value in bestpupil[3].items():
            if value != 0:
                pupil = prysm.FringeZernike(dia=10, norm=False,
                                            wavelength=bestpupil[2],
                                            opd_unit="um",
                                            samples=samples,
                                            **{key: value})
                pupil = mask_pupil(pupil, p['base_fstop'], p['fstop'], engine=me)
                slice = pupil.slice_x[1]
                rms = (slice[me.isfinite(slice)] ** 2).mean() ** 0.5
                plt.plot(cp.asnumpy(slice), label="{} : {:.3f} λRMS".format(key, rms / conf.BASE_WAVELENGTH))
        slice_ = bestpupil[1].slice_x[1]
        rms = (slice_[me.isfinite(slice_)] ** 2).mean() ** 0.5
        plt.plot(slice_, label="All : {:.3f} λRMS".format(rms / conf.BASE_WAVELENGTH))
        pupil = prysm.FringeZernike(z4=z4,
                                    dia=10, norm=False,
                                  wavelength=bestpupil[2],
                                  opd_unit="um",
                                  samples=samples,
                                  **bestpupil[3])

        pupil = mask_pupil(pupil, p['base_fstop'], p['fstop'])
        slice_ = pupil.slice_x[1]
        rms = (slice_[me.isfinite(slice_)] ** 2).mean() ** 0.5
        plt.plot(slice_, '--', label="ZeroZ4 : {:.3f} λRMS".format(rms / conf.BASE_WAVELENGTH), color='black')
        plt.legend()
        plt.show()
        mono_psf = prysm.PSF.from_pupil(pupil, efl=p['base_fstop']*10)
        mono_psf.plot2d(axlim=8)
        plt.show()

        slicedv = np.abs(np.diff(sliceavg[me.isfinite(slice_)]))
        peakiness = slicedv.max() / slicedv.mean()

        plt.plot(np.diff(slice_), label="Pupil slice derivative at best focus {:.3f}".format(peakiness))
        plt.legend()
        plt.show()

    return tr

def _try_wavefront_prysmref(s: TestSettings):
    t = time.time()

    tr = TestResults()
    tr.copy_important_settings(s)

    if s.dummy:
        return tr

    if s.p['fstop'] < s.p['base_fstop']:
        raise ValueError("Base_fstop must be wider (lower) than fstop {} < {}".format(s.p['fstop'], s.p['base_fstop']))
    pupilslices = []

    if not s.is_valid:
        s = get_processing_details(s)

    use_cuda = False

    if prysm.config.backend != np:
        prysm.config.backend = np
        prysm.config.precision = conf.PRECISION

    tr.used_cuda = use_cuda

    realdtype = "float32" if conf.PRECISION == 32 else "float64"
    complexdtype = "complex64" if conf.PRECISION == 32 else "complex128"

    me, mestr = (cp, 'np') if use_cuda else (np, 'np')

    eval_wavelengths = [BASE_WAVELENGTH] if s.mono else MODEL_WVLS

    # eval_wavelengths = [BASE_WAVELENGTH] if s.mono else [0.5, 0.9]

    # s.p['loca'] *= 0.1
    # s.p['loca1'] *= 0.1
    # s.p['spca2'] *= 0.1
    # s.p['spca'] *= 0.1

    psf_lst = []

    # s.fftsize = s.phasesamples * 4

    prysm_path_q = max(1, s.fftsize / s.phasesamples)

    samplelst = []
    print(s.id_or_hash)

    polychromatic_weights = me.array([float(photopic_fn(wv * 1e3) * d50_interpolator(wv)) for wv in eval_wavelengths])
    # polychromatic_weights = me.array([0.5, 0.5])
    t_misc = 0
    t_pupils = 0
    t_get_phases = 0
    t_get_fcns = 0
    t_pads = 0
    t_ffts = 0
    t_cudasyncs = 0
    t_affines = 0
    t_mtfs = 0
    t_init = time.time() - t
    t = time.time()

    mask = mask_pupil(s, engine=me, dtype=realdtype)

    t_maskmaking = time.time() - t

    bestpupil = None

    zkwargs = {}
    for key, value in s.p.items():
        if key.lower().startswith('z') and key[1].isdigit():
            if key.lower() != "z4" and key.lower() != "z9":
                zkwargs[key] = value * conf.BASE_WAVELENGTH

    for wvl_num, (model_wvl, polych_weight) in enumerate(zip(eval_wavelengths, polychromatic_weights)):
        t = time.time()
        rel_wv = model_wvl / BASE_WAVELENGTH

        z4 = get_z4(s.defocus, s.p, model_wvl)
        z9 = get_z9(s.p, model_wvl)

        samplelst.append(s.phasesamples)
        t_misc += time.time() - t
        t = time.time()
        pupil = prysm.FringeZernike(z4=z4, z9=z9, dia=10, norm=False,
                                    wavelength=model_wvl,
                                    opd_unit="um",
                                    mask_target='fcn',
                                    mask=mask,
                                    samples=s.phasesamples, **zkwargs)
        t_pupils += time.time() - t
        t = time.time()

        t_get_phases += time.time() - t
        t = time.time()

        t_get_phases += time.time() - t

        t = time.time()
        mono_psf = prysm.PSF.from_pupil(pupil, efl=s.p['base_fstop'] * 10, Q=prysm_path_q, norm='radiometric')
        # plt.plot(np.log10(mono_psf.data[int(mono_psf.data.shape[0]/2), :]))
        # mono_psf.plot2d()
        # plt.show()

        t_ffts += time.time() - t

        psf_lst.append(mono_psf)

        pupilslices.append(pupil.slice_x[1])

        metric = np.abs(rel_wv - 1)
        if bestpupil is None or metric < bestpupil[0]:
            bestpupil = metric, pupil, model_wvl, zkwargs

    t = time.time()

    mixedpsf = prysm.PSF.polychromatic(psf_lst, me.array(polychromatic_weights))
    tr.psf = mixedpsf
    # print(s.id_or_hash)
    # tr2 = _try_wavefront_prysmref(s)
    # mixedpsf.plot2d(axlim=18)
    # ax = plt.gca()
    # ax.set_title("A")
    # plt.show()
    # tr.otf = np.ones_like(SPACIAL_FREQS), np.ones_like(SPACIAL_FREQS)

    t_ffts += time.time() - t

    t = time.time()
    freqs = SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3
    mtf = prysm.MTF.from_psf(mixedpsf)
    tr._prysm_mtf = mtf.sag, mtf.tan
    tr.otf = mtf.exact_sag(freqs), mtf.exact_tan(freqs)
    # tr._prysm_otf = tr.otf
    if 0:
        centre = s.fftsize // 2

        otf = me.fft.fftshift(me.fft.fft2(me.fft.ifftshift(mixedpsf.data)))

        f_units = prysm.fttools.forward_ft_unit(mixedpsf.sample_spacing, len(mixedpsf.x))
        otf /= np.abs(otf[centre, centre])
        sag_mod = otf[centre, centre:]
        tan_mod = otf[centre:, centre]
        f_units = f_units[centre:]

        interpfn = interpolate.InterpolatedUnivariateSpline(f_units, np.real(sag_mod), k=1)
        sagmtf = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)
        interpfn = interpolate.InterpolatedUnivariateSpline(f_units, np.real(tan_mod), k=1)
        tanmtf = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)
        interpfn = interpolate.InterpolatedUnivariateSpline(f_units, np.imag(sag_mod), k=1)
        sagmtf_i = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)
        interpfn = interpolate.InterpolatedUnivariateSpline(f_units, np.imag(tan_mod), k=1)
        tanmtf_i = interpfn(SPACIAL_FREQS / DEFAULT_PIXEL_SIZE * 1e-3)

        tr.otf = sagmtf + 1j * sagmtf_i, tanmtf + 1j * tanmtf_i

    t_mtfs += time.time() - t

    timings = dict(t_init=t_init,
                   t_maskmaking=t_maskmaking,
                   t_pupils=t_pupils,
                   t_get_phases=t_get_phases,
                   t_get_fcns=t_get_fcns,
                   t_pads=t_pads,
                   t_ffts=t_ffts,
                   t_cudasyncs=t_cudasyncs,
                   t_affines=t_affines,
                   t_mtfs=t_mtfs,
                   t_misc=t_misc)

    tr.timings = timings

    sliceavg = me.average(me.array(pupilslices, dtype='float64'), axis=0, weights=polychromatic_weights)

    slice_ = bestpupil[1].slice_x[1]
    slicedv = me.abs(me.diff(sliceavg[me.isfinite(sliceavg)]))
    if slicedv.sum() != 0:
        peakiness = slicedv.max() / slicedv.mean()
    else:
        peakiness = 999.0
    # strehl = float(bestpupil[1].strehl)
    strehl = 1

    if s.plot:
        # Plot each Z phase separately
        for key, value in bestpupil[3].items():
            if value != 0:
                pupil = prysm.FringeZernike(dia=10, norm=False,
                                            wavelength=bestpupil[2],
                                            opd_unit="um",
                                            samples=samples,
                                            **{key: value})
                pupil = mask_pupil(pupil, p['base_fstop'], p['fstop'], engine=me)
                slice = pupil.slice_x[1]
                rms = (slice[me.isfinite(slice)] ** 2).mean() ** 0.5
                plt.plot(cp.asnumpy(slice), label="{} : {:.3f} λRMS".format(key, rms / conf.BASE_WAVELENGTH))
        slice_ = bestpupil[1].slice_x[1]
        rms = (slice_[me.isfinite(slice_)] ** 2).mean() ** 0.5
        plt.plot(slice_, label="All : {:.3f} λRMS".format(rms / conf.BASE_WAVELENGTH))
        pupil = prysm.FringeZernike(z4=z4,
                                    dia=10, norm=False,
                                    wavelength=bestpupil[2],
                                    opd_unit="um",
                                    samples=samples,
                                    **bestpupil[3])

        pupil = mask_pupil(pupil, p['base_fstop'], p['fstop'])
        slice_ = pupil.slice_x[1]
        rms = (slice_[me.isfinite(slice_)] ** 2).mean() ** 0.5
        plt.plot(slice_, '--', label="ZeroZ4 : {:.3f} λRMS".format(rms / conf.BASE_WAVELENGTH), color='black')
        plt.legend()
        plt.show()
        mono_psf = prysm.PSF.from_pupil(pupil, efl=p['base_fstop']*10)
        mono_psf.plot2d(axlim=8)
        plt.show()

        slicedv = np.abs(np.diff(sliceavg[me.isfinite(slice_)]))
        peakiness = slicedv.max() / slicedv.mean()

        plt.plot(np.diff(slice_), label="Pupil slice derivative at best focus {:.3f}".format(peakiness))
        plt.legend()
        plt.show()

    return tr

tempcache = {}

def mask_pupil(s: TestSettings, engine=np, dtype="float64", plot=False):
    # s.p['v_rad'] = 1.0
    hashtuple = (s.p['base_fstop'],
                 s.p['fstop'],
                 s.x_loc,
                 s.y_loc,
                 s.phasesamples,
                 s.p.get('a', 0.0),
                 s.p.get('b', 0.0),
                 s.p.get('v_scr', 1.0),
                 s.p.get('v_rad', 1.0),
                 # s.p.get('v_x', 0.0),
                 # s.p.get('v_y', 0.0),
                 s.p.get('squariness', 0.5),
                 "np" if engine is np else "cp",
                 dtype)

    hash = hashtuple.__hash__()

    for cachehash, mask in mask_cache:
        if cachehash == hash:
            return mask

    smoothfactor = s.phasesamples / 1.5

    me = engine

    aperture_stop_norm_radius = s.p['base_fstop'] / s.p['fstop']

    na = 1 / (2.0 * s.p['base_fstop'])
    onaxis_peripheral_ray_angle = me.arcsin(na, dtype=dtype)

    pupil_radius_mm = me.tan(onaxis_peripheral_ray_angle, dtype=dtype) * s.default_exit_pupil_position_mm

    x_displacement_mm = (s.x_loc - IMAGE_WIDTH / 2) * DEFAULT_PIXEL_SIZE * 1e3
    y_displacement_mm = (s.y_loc - IMAGE_HEIGHT / 2) * DEFAULT_PIXEL_SIZE * 1e3

    # angle = s.p.get('v_angle', 0)
    magnitude = (x_displacement_mm ** 2 + y_displacement_mm ** 2) ** 0.5

    if s.fix_pupil_rotation:
        x_displacement_mm = -magnitude
        y_displacement_mm = 0

    x_displacement_mm_min = -x_displacement_mm - pupil_radius_mm
    x_displacement_mm_max = -x_displacement_mm + pupil_radius_mm
    y_displacement_mm_min = -y_displacement_mm - pupil_radius_mm
    y_displacement_mm_max = -y_displacement_mm + pupil_radius_mm

    x = me.linspace(x_displacement_mm_min, x_displacement_mm_max, s.phasesamples, dtype=dtype)
    y = me.linspace(y_displacement_mm_min, y_displacement_mm_max, s.phasesamples, dtype=dtype)
    gridx, gridy = me.meshgrid(x, y)
    displacement_grid = (gridx**2 + gridy**2) ** 0.5
    squariness = (2**0.5 - displacement_grid / me.maximum(abs(gridx), abs(gridy))) ** 2
    pixel_angle_grid = me.arctan(displacement_grid / s.default_exit_pupil_position_mm *
                                 (1.0 + squariness * s.p.get('squariness', 0.5)), dtype=dtype)

    normarr = me.linspace(-1, 1, s.phasesamples, dtype=dtype)
    gridx, gridy = me.meshgrid(normarr, normarr)
    pupil_norm_radius_grid = (gridx ** 2 + gridy ** 2) ** 0.5

    stopmask = np.clip( (aperture_stop_norm_radius - pupil_norm_radius_grid) * smoothfactor + 0.5, 0, 1)

    a = s.p.get('a', 1.0)
    b = s.p.get('b', 1.0)

    if not s.pixel_vignetting:
        mask = stopmask
    else:
        coeff_4 = -18.73 * a
        corff_6 = 485 * b
        square_grid = 1.0 / (1.0 + (pixel_angle_grid**4 * coeff_4 + pixel_angle_grid ** 6 * corff_6))
        mask = stopmask * square_grid

    if s.lens_vignetting:
        # Lens vignette
        normarr = me.linspace(-1, 1, s.phasesamples, dtype=dtype)
        image_circle_modifier = s.p.get('v_slr', 1.0) * 0.6
        gridx, gridy = me.meshgrid(normarr - x_displacement_mm / SENSOR_WIDTH * 1e-3 * image_circle_modifier,
                                   normarr - y_displacement_mm / SENSOR_WIDTH * 1e-3 * image_circle_modifier)
        vignette_radius_grid = (gridx ** 2 + gridy ** 2) ** 0.5
        vignette_crop_circle_radius = s.p.get('v_rad', 1.0)
        vignette_mask = vignette_radius_grid < vignette_crop_circle_radius
        vignette_mask = np.clip((vignette_crop_circle_radius - vignette_radius_grid) * smoothfactor + 0.5, 0, 1)
        mask *= vignette_mask

        normarr = me.linspace(-1, 1, s.phasesamples, dtype=dtype)
        image_circle_modifier = s.p.get('v_slr', 1.0) * 0.6
        gridx, gridy = me.meshgrid(normarr + x_displacement_mm / SENSOR_WIDTH * 1e-3 * image_circle_modifier,
                                   normarr + y_displacement_mm / SENSOR_WIDTH * 1e-3 * image_circle_modifier)
        vignette_radius_grid = (gridx ** 2 + gridy ** 2) ** 0.5
        vignette_crop_circle_radius = s.p.get('v_rad', 1.0) * 1.0
        vignette_mask = vignette_radius_grid < vignette_crop_circle_radius
        vignette_mask = np.clip((vignette_crop_circle_radius - vignette_radius_grid) * smoothfactor + 0.5, 0, 1)
        mask *= vignette_mask

    if plot or s.id_or_hash == -1:
        if engine is cp:
            print(square_grid)
            plt.imshow(cp.asnumpy(mask))
            plt.colorbar()
            plt.show()
        else:
            print(square_grid)
            plt.imshow(mask)
            plt.colorbar()
            plt.show()
    return mask

def plot_pixel_vignetting_loss():
    fstops = 2.0 ** np.linspace(0.0, 1.5, 6)
    test_fstops = (1, 2**0.16667, 1.222, 2**0.5, 1.4 * 2**0.166667, 2, 2 * 2 ** 0.166667, 2 * 2**0.5)
    benefits_exp = (1.93,1.90, 1.81, 1.70, 1.5, 0.93, 0.62, 0)

    s = TestSettings(0, dict(base_fstop=1.0, fstop=2.0 * 2**0.5))
    s.phasesamples = 256
    s.pixel_vignetting = False
    baseline = mask_pupil(s, np).mean()

    s.pixel_vignetting = True

    if "optimise" and 0:
        def callable(x):
            error = 0
            for testfstop, benefit_exp in zip(test_fstops, benefits_exp):
                s.p = dict(base_fstop=1, fstop=testfstop)
                s.p['a'], s.p['b'] = x
                benefit = np.log2(mask_pupil(s, np).mean() / baseline)
                print(x[0], x[1], testfstop, benefit)
                error += (benefit - benefit_exp) ** 2
            print()
            return error

        opt = optimize.minimize(callable, (0, 0))  # ,bounds=((-20,20), (-30, 30)
        a, b = opt.x
    else:
        a, b = 0, 0

    plot_fstops = 2**np.linspace(0, 1.5, 6)
    benefits = []
    benefits_stop = []
    for testfstop in plot_fstops:
        s.p = dict(base_fstop=testfstop, fstop=testfstop)
        s.p['a'], s.p['b'] = a, b
        benefit = np.log2(mask_pupil(s, np, plot=True).mean()/(testfstop**2) / baseline)
        benefits.append(benefit)
        s.p = dict(base_fstop=1, fstop=testfstop)
        s.p['a'], s.p['b'] = a, b
        benefit = np.log2(mask_pupil(s, np, plot=True).mean() / baseline)
        benefits_stop.append(benefit)

    plt.plot(plot_fstops, benefits)
    plt.plot(plot_fstops, benefits_stop)
    plt.plot(test_fstops, benefits_exp)
    plt.plot()
    plt.show()
    exit()

def plot_lens_vignetting_loss(base_fstop=1.4):
    fstops = 2.0 ** np.linspace(0.0, 2, 4)
    for stop in fstops:
        s = TestSettings(0, dict(base_fstop=base_fstop, fstop=stop * base_fstop))
        s.phasesamples = 128
        s.pixel_vignetting = True
        s.lens_vignetting = True
        s.p['v_mag'] = 0.8
        s.p['v_rad'] = 1.3
        s.p['v_x'] = -0.8
        s.p['v_y'] = -0.8
        baseline = mask_pupil(s, np).mean()
        heights = np.linspace(0, 1, 16)
        losses = []
        s.x_loc = 0
        s.y_loc = 0
        mask_pupil(s, np, plot=True)
        for height in heights:
            s.x_loc = 3000 + height * IMAGE_WIDTH / 2
            s.y_loc = 2000 + height * IMAGE_HEIGHT / 2
            # losses.append(np.log2(mask_pupil(s, np).mean() / baseline))
            losses.append(mask_pupil(s, np).mean() / baseline)
        plt.plot(heights, losses)
    plt.show()

# from lentil import wavefront_test
# s = TestSettings(0, dict(base_fstop=2.6, fstop=2.6))
# s.phasesamples = 256
# s.return_type = RETURN_MTF
# s.p['tca_slr'] = 0
#
# s.p['v_slr'] =2.6
# s.x_loc = 5600
# s.y_loc = 3800
# s.x_loc = 6000-5600
# s.y_loc = 4000-3800
# s.p['z5'] = -1
# s.p['df_offset'] = 0
# s.p['df_step'] = 1
# s.default_exit_pupil_position_mm = 1000
# s.p['ellip'] = 0
# s.allow_cuda = False
# s.p['loca'] = 0.0000001
# s.defocus = 0
# s.mono = True

# mask_pupil(s, plot=True)
# tr = try_wavefront(s)
# plt.plot(tr.otf[0], label='sag')
# plt.plot(tr.otf[1], label='tan')
# plt.legend()
# print(tr.otf[0])
# psf = tr.psf
# psf.plot2d(axlim=30)
# plt.show()
# exit()y

pool = multiprocessing.Pool(processes=8)

def testmul(samples, total_loops=2*5, Q=2):
    s = TestSettings(0, dict(fstop=1.4, base_fstop=1.4))
    s.phasesamples = 256
    s.fftsize = 512
    s.allow_cuda = False
    pool.starmap(try_wavefront, [(s, )] * 16 * 5)


#!/usr/bin/env python

import sys
import numpy as np
from libs import base
from libs import utils
import libs.cadnano_utils as cu
import re
import os
import pickle

DEBUG = 0
DIST_HEXAGONAL = 2.55  # distance between centres of virtual helices (hexagonal array)
DIST_SQUARE = 2.60  # distance between centres of virtual helices (square array)
BOX_FACTOR = 2  # factor by which to expand the box (linear dimension)

class vh_nodes(object):

    def __init__(self):
        self.begin = []
        self.end = []

    def __str__(self):
        return str([self.begin, self.end])

    def add_begin(self, begin_index):
        if begin_index not in self.begin:
            self.begin.append(begin_index)

    def add_end(self, end_index):
        if end_index not in self.end:
            self.end.append(end_index)

def vhelix_rotation_origami_sq(direction, perp):
    R = utils.get_rotation_matrix(direction, np.pi * 15. / 180)
    return np.dot(R, perp)

def vhelix_rotation_origami_he(direction, perp):
    R = utils.get_rotation_matrix(direction, np.pi * 160. / 180)
    return np.dot(R, perp)

def insert_loop_skip(strands, start_pos, direction, perp, rot, helix_angles, vhelix, nodes, use_seq, seqs):
    # return a double strand which is a copy of the double strand in the first argument, but with skips and loops

    # strand is generated right to left i.e. opposite direction to even vhelix
    length_change = []
    length_change_total = 0
    new_nodes = vh_nodes()
    new_angle = []
    helix_angles_new = np.copy(helix_angles)

    if vhelix.num % 2 == 1:
        reverse_nodes = vh_nodes()
        reverse_nodes.begin = list(reversed(nodes.begin))
        reverse_nodes.end = list(reversed(nodes.end))

    for i in range(len(nodes.begin)):
        # ltr: left to right; looking at the strand left to right (low to high square index), the beginning/end of the effective strand is here (before skips/loops)
        # gs: generated strand; the index of the nucleotide (BEFORE skips/loops are applied) on the generated strand corresponding to the beginning/end of the effective strand
        if vhelix.num % 2 == 0:
            begin_ltr = nodes.begin[i]
            end_ltr = nodes.end[i]
            begin_gs = nodes.begin[i]
            end_gs = nodes.end[i]
        else:
            begin_ltr = reverse_nodes.end[i]
            end_ltr = reverse_nodes.begin[i]
            begin_gs = vhelix.len - reverse_nodes.begin[i] - 1
            end_gs = vhelix.len - reverse_nodes.end[i] - 1

        # check for zero length effective strand
        if end_gs - begin_gs != 0:
            # get length change for this effective strand
            length_change.append(0)
            for j in vhelix.skip[begin_ltr:end_ltr + 1]:
                length_change[i] -= int(j)
            for j in vhelix.loop[begin_ltr:end_ltr + 1]:
                length_change[i] += int(j)
            # get new pitch angles for this effective strand
            new_angle.append(sum(helix_angles[begin_gs:end_gs]) / (end_gs - begin_gs + length_change[i]))
            helix_angles_new[begin_gs:end_gs] = new_angle[i]
            # adjust beginning/end indices according to length change
            begin_gs += length_change_total
            end_gs += length_change_total + length_change[i]
            new_nodes.add_begin(begin_gs)
            new_nodes.add_end(end_gs)  # begin_gs > end_gs.....
        else:
            length_change.append(0)
            new_angle.append(sum(helix_angles) / len(helix_angles))  # append an average angle
            new_nodes.add_begin(begin_gs)
            new_nodes.add_end(end_gs)
        length_change_total += length_change[i]

    # adjust the new helix angle array according to skips/loops
    deleted = 0
    inserted = 0
    deleted_this_iteration = 0
    inserted_this_iteration = 0
    for i in range(len(nodes.begin)):
        deleted += deleted_this_iteration
        inserted += inserted_this_iteration
        deleted_this_iteration = 0
        inserted_this_iteration = 0
        if vhelix.num % 2 == 0:
            begin_ltr = nodes.begin[i]
            end_ltr = nodes.end[i]
            begin_gs = nodes.begin[i]
            end_gs = nodes.end[i]
        else:
            begin_ltr = reverse_nodes.end[i]
            end_ltr = reverse_nodes.begin[i]
            begin_gs = vhelix.len - reverse_nodes.begin[i] - 1
            end_gs = vhelix.len - reverse_nodes.end[i] - 1
        for j in vhelix.skip[begin_ltr:end_ltr + 1]:
            if j == 1:
                helix_angles_new = np.delete(helix_angles_new, begin_gs - deleted + inserted)
                deleted_this_iteration += 1
        for j in vhelix.loop[begin_ltr:end_ltr + 1]:
            for _ in range(j):
                helix_angles_new = np.insert(helix_angles_new, begin_gs - deleted + inserted, new_angle[i])
                inserted_this_iteration += 1

    g = cu.StrandGenerator()
    new_strands = g.generate_or_sq(len(helix_angles_new) + 1, start_pos=start_pos, direction=direction, perp=perp, double=True, rot=rot, angle=helix_angles_new, length_change=length_change, region_begin=new_nodes.begin, region_end=new_nodes.end)
    if use_seq:
        try:
            sequence = [x for x in seqs[vhelix.cad_index]]
        except IndexError:
            base.Logger.die("sequence file contains too few rows compared to the number of virtual helices in the cadnano file, dying")
        if vhelix.num % 2 == 1:
            sequence.reverse()
        if new_strands[0].get_length() != len(sequence):
            base.Logger.log("Cannot change sequence: lengths don't match; virtual helix %s, sequence length %s, virtual helix length %s - are skips/loops accounted for?" % (vhelix.num, len(sequence), new_strands[0].get_length()), base.Logger.WARNING)
        else:
            new_strands[0].set_sequence(sequence)
        sequence2 = [3 - s for s in sequence]
        sequence2.reverse()
        if new_strands[0].get_length() != len(sequence):
            base.Logger.log("Cannot change sequence: lengths don't match; virtual helix %s, sequence length %s, virtual helix length %s - are skips/loops accounted for?" % (vhelix.num, len(sequence), new_strands[0].get_length()), base.Logger.WARNING)
        else:
            new_strands[1].set_sequence(sequence2)

    return new_strands

def add_slice(current_system, vhelix, begin, end, nodes, strands, pos, direction, perp, rot, helix_angles, strand_type, use_seq, seqs):
    # add a slice of the virtual helix to the slice system, taking into account skips and loops
    length_change_begin = 0
    length_change_end = 0

    if (vhelix.num % 2 + strand_type) % 2 == 0:  # strand and even num or staple and odd num
        for i in vhelix.skip[:begin]:
            length_change_begin -= int(i)
        for i in vhelix.skip[:end + 1]:
            length_change_end -= int(i)

        for i in vhelix.loop[:begin]:
            length_change_begin += int(i)
        for i in vhelix.loop[:end + 1]:
            length_change_end += int(i)

        begin_slice = begin + length_change_begin
        end_slice = end + 1 + length_change_end

    else:
        for i in vhelix.skip[end:]:
            length_change_end -= int(i)
        for i in vhelix.skip[begin + 1:]:
            length_change_begin -= int(i)

        for i in vhelix.loop[end:]:
            length_change_end += int(i)
        for i in vhelix.loop[begin + 1:]:
            length_change_begin += int(i)

        begin_slice = vhelix.len - begin - 1 + length_change_begin
        end_slice = vhelix.len - end + length_change_end

    new_strands = insert_loop_skip(strands, pos, direction, perp, rot, helix_angles, vhelix, nodes, use_seq, seqs)
    current_system.add_strand(new_strands[strand_type].get_slice(begin_slice, end_slice), check_overlap=False)
    return current_system

def add_slice_nupack(vhelix, strand_number, begin_helix, end_helix, index_lookup, strand_type):
    length_change = 0
    skips = 0
    loops = 0
    if (vhelix.num % 2 + strand_type) % 2 == 0 :  # strand and even num or staple and odd num
        for i in vhelix.skip[begin_helix:end_helix + 1]:
            length_change -= int(i)

        for i in vhelix.loop[begin_helix:end_helix + 1]:
            length_change += int(i)

    else:
        for i in vhelix.skip[end_helix:begin_helix + 1]:
            length_change -= int(i)

        for i in vhelix.loop[end_helix:begin_helix + 1]:
            length_change += int(i)

    if (strand_type + vhelix.num % 2) % 2 == 0:
        iter_length = end_helix - begin_helix + 1 + length_change
    else:
        iter_length = begin_helix + 1 - end_helix + length_change

    nucleotide = 0
    while nucleotide < iter_length:
        if (strand_type + vhelix.num % 2) % 2 == 0:
            vhelix_base = nucleotide + begin_helix + skips - loops
        else:
            vhelix_base = begin_helix - nucleotide - skips + loops
        if vhelix.skip[vhelix_base] != 1:
            if (strand_type + vhelix.num % 2) % 2 == 0:
                add_nuc = [nucleotide + x for x in range(vhelix.loop[vhelix_base] + 1)]
            else:
                add_nuc = [nucleotide + x for x in range(vhelix.loop[vhelix_base] + 1)[::-1]]
            if strand_type == 0:
                index_lookup[(vhelix.num, vhelix_base)] = [strand_number, nucleotide]
            elif strand_type == 1:
                index_lookup[(strand_number, nucleotide)] = [vhelix.num, vhelix_base]
            elif strand_type == 2:
                index_lookup.add_scaf(vhelix.num, vhelix_base, strand_number, add_nuc)
            elif strand_type == 3:
                index_lookup.add_stap(vhelix.num, vhelix_base, strand_number, add_nuc)
            nucleotide += 1 + vhelix.loop[vhelix_base]
            loops += vhelix.loop[vhelix_base]
        else:
            if strand_type == 2:
                index_lookup.add_scaf(vhelix.num, vhelix_base, strand_number, [])
            elif strand_type == 3:
                index_lookup.add_stap(vhelix.num, vhelix_base, strand_number, [])
            skips += 1

    return index_lookup

def build_nodes(vh):
    # returns a vh_nodes object which contains the beginning and end square indices of each effective strand. Effective strands
    # on a given vhelix are a combination of information on staple and scaffold strands. Together they tell us which
    # nucleotides need to be held constant when we alter the angles/base-base distances for a section of nucleotides along a final strand
    nodes = vh_nodes()
    if vh.num % 2 == 0:
        direction = 1
    else:
        direction = -1
    for i in range(len(vh.scaf)):
        # need to consider what happens when I add an index to the node list that doesn't fall within the range of square indices in the vhelix
        previd = i - 1 * direction
        nextid = i + 1 * direction
        if previd in range(len(vh.scaf)):
            prev = vh.scaf[previd].type(vh, previd)
            prev_stap = vh.stap[previd].type(vh, previd)
        else:
            prev = False
            prev_stap = False
        if nextid in range(len(vh.scaf)):
            next_ = vh.scaf[nextid].type(vh, nextid)
            next_stap = vh.stap[nextid].type(vh, nextid)
        else:
            next_ = False
            next_stap = False
        # now build the effective strand vh_nodes object
        if not ((prev == prev_stap and (prev == 'begin' or prev == 'end')) or (next_ == next_stap and (next_ == 'begin' or next_ == 'end'))):
            if vh.scaf[i].type(vh, i) == 'empty':
                if vh.stap[i].type(vh, i) == 'begin':
                    nodes.add_end(i)
                elif vh.stap[i].type(vh, i) == 'end':
                    nodes.add_begin(i)
            elif vh.scaf[i].type(vh, i) == 'begin':
                if vh.stap[i].type(vh, i) == 'empty':
                    nodes.add_begin(i)
                elif vh.stap[i].type(vh, i) == 'continue':
                    nodes.add_begin(i)
                    nodes.add_end(i - 1 * direction)
                elif vh.stap[i].type(vh, i) == 'begin':
                    nodes.add_begin(i + 1 * direction)
                    nodes.add_end(i - 1 * direction)
                elif vh.stap[i].type(vh, i) == 'end':
                    nodes.add_begin(i)
            elif vh.scaf[i].type(vh, i) == 'end':
                if vh.stap[i].type(vh, i) == 'empty':
                    nodes.add_end(i)
                elif vh.stap[i].type(vh, i) == 'continue':
                    nodes.add_begin(i + 1 * direction)
                    nodes.add_end(i)
                elif vh.stap[i].type(vh, i) == 'begin':
                    nodes.add_end(i)
                elif vh.stap[i].type(vh, i) == 'end':
                    nodes.add_begin(i + 1 * direction)
                    nodes.add_end(i - 1 * direction)
            elif vh.scaf[i].type(vh, i) == 'continue':
                if vh.stap[i].type(vh, i) == 'begin':
                    nodes.add_begin(i + 1 * direction)
                    nodes.add_end(i)
                elif vh.stap[i].type(vh, i) == 'end':
                    nodes.add_begin(i)
                    nodes.add_end(i - 1 * direction)

    return nodes

def generate_vhelices_origami_sq(vhelix_direction, vhelix_perp, h, sequence_file, single_strand_system, vhelix_counter):
    g = cu.StrandGenerator()
    # generate helix angles
    helix_angles = np.zeros(h.len - 1, dtype=float)
    # hard upper limit on pitch angle seems to be between 54.5 and 55 degrees
    for i in range(len(helix_angles)):
        modi = i % 32
        if modi < 2:
            helix_angles[i] = 28 * np.pi / 180
        elif modi == 2:
            helix_angles[i] = 36 * np.pi / 180
        elif modi == 3:
            helix_angles[i] = 54.375 * np.pi / 180
        elif modi == 4:
            helix_angles[i] = 37 * np.pi / 180
        elif modi in (5, 6):
            helix_angles[i] = 27.6666666666666 * np.pi / 180
        elif modi == 7:
            helix_angles[i] = 30.6666666666666 * np.pi / 180
        elif modi in (8, 9):
            helix_angles[i] = 29.3333333333 * np.pi / 180
        elif modi == 10:
            helix_angles[i] = 34.3333333333 * np.pi / 180
        elif modi == 11:
            helix_angles[i] = 54.5 * np.pi / 180
        elif modi in (12, 13):
            helix_angles[i] = (28.91666666666 * np.pi / 180)  # + 0.25) * np.pi/180
        elif modi in (14, 15, 16, 17):
            helix_angles[i] = 31.16666666666 * np.pi / 180
        elif modi == 18:
            helix_angles[i] = 35.5 * np.pi / 180
        elif modi == 19:
            helix_angles[i] = 52 * np.pi / 180
        elif modi == 20:
            helix_angles[i] = 35.5 * np.pi / 180
        elif modi in (21, 22):
            helix_angles[i] = 27.5 * np.pi / 180
        elif modi == 23:
            helix_angles[i] = 35.5 * np.pi / 180
        elif modi >= 24 and modi < 27:
            helix_angles[i] = 30 * np.pi / 180
        elif modi == 27:
            helix_angles[i] = 52 * np.pi / 180
        elif modi == 28:
            helix_angles[i] = 35.5 * np.pi / 180
        else:
            helix_angles[i] = 30.91666666666 * (np.pi / 180)

    # make sure the helices are periodic in 32 bases
    total_sum = 0
    for i in range(31):
        total_sum += helix_angles[i]

    for i in range(len(helix_angles)):
        if i % 32 == 31:
            helix_angles[i] = 1080 * np.pi / 180 - total_sum

    # make the virtual helices
    if h.num % 2 == 0:
        pos = np.array([h.col * DIST_SQUARE, h.row * DIST_SQUARE, 0])
        direction = vhelix_direction
        perp = vhelix_perp
        rot = 0.
        angles = helix_angles
        strands = g.generate_or_sq(h.len, start_pos=pos, direction=direction, perp=perp, double=True, rot=rot, angle=angles)

    else:
        pos = np.array([h.col * DIST_SQUARE, h.row * DIST_SQUARE, (h.len - 1) * base.BASE_BASE])
        direction = -vhelix_direction
        perp = -vhelix_perp
        rot = -np.sum(helix_angles) % (2 * np.pi)
        angles = np.flipud(helix_angles)
        strands = g.generate_or_sq(h.len, start_pos=pos, direction=direction, perp=perp, double=True, rot=rot, angle=angles)

    return (strands[0], strands[1]), helix_angles, pos, rot, direction, perp

def generate_vhelices_origami_he(vhelix_direction, vhelix_perp, h, sequence_file, single_strand_system, vhelix_counter):
    g = cu.StrandGenerator()
    # generate helix angles
    helix_angles = np.zeros(h.len - 1, dtype=float)

    for i in range(len(helix_angles)):
        modi = i % 21
        if modi == 0:
            helix_angles[i] = 32.571 * np.pi / 180
        elif modi == 1:
            helix_angles[i] = 36 * np.pi / 180
        elif modi in (1, 2, 3):
            helix_angles[i] = 42 * np.pi / 180
        elif modi in (5, 6, 7):
            helix_angles[i] = 29.143 * np.pi / 180
        elif modi == 8:
            helix_angles[i] = 32 * np.pi / 180
        elif modi in (9, 10):
            helix_angles[i] = 44 * np.pi / 180
        elif modi in (12, 13, 14):
            helix_angles[i] = 28.571 * np.pi / 180
        elif modi in (16, 17):
            helix_angles[i] = 41.5 * np.pi / 180
        elif modi in (19, 20):
            helix_angles[i] = 28.476 * np.pi / 180
        else:
            helix_angles[i] = 720. / 21 * (np.pi / 180.)

    # make sure it's periodic
    total_sum = 0
    for i in range(20):
        total_sum += helix_angles[i]

    for i in range(len(helix_angles)):
        if i % 21 == 20:
            helix_angles[i] = 720. * np.pi / 180 - total_sum

    # make the virtual helices
    if h.num % 2 == 0:
        pos = np.array([h.col * np.sqrt(3) * DIST_HEXAGONAL / 2, h.row * 3 * DIST_HEXAGONAL / 2, 0])
        direction = vhelix_direction
        perp = vhelix_perp
        rot = 0.
        strands = g.generate_or_sq(h.len, start_pos=pos, direction=direction, perp=perp, double=True, rot=rot, angle=helix_angles)

    else:
        pos = np.array([h.col * np.sqrt(3) * DIST_HEXAGONAL / 2, h.row * 3 * DIST_HEXAGONAL / 2 + DIST_HEXAGONAL / 2, (h.len - 1) * base.BASE_BASE])
        direction = -vhelix_direction
        perp = -vhelix_perp
        if base.MM_GROOVING:
            rot = -np.sum(helix_angles) % (2 * np.pi) - 0.07
        else:
            rot = -np.sum(helix_angles) % (2 * np.pi)
        angles = np.flipud(helix_angles)
        strands = g.generate_or_sq(h.len, start_pos=pos, direction=direction, perp=perp, double=True, rot=rot, angle=angles)

    return (strands[0], strands[1]), helix_angles, pos, rot, direction, perp

# cadnano object structure
class vstrands (object):

    def __init__(self):
        self.vhelices = []

    def add_vhelix(self, toadd):
        self.vhelices.append(toadd)

    def bbox(self):
        rows = []
        cols = []
        lens = []
        for h in self.vhelices:
            rows.append(h.row)
            cols.append(h.col)
            lens.append(len(h.stap))

        dr = DIST_SQUARE * (max(rows) - min(rows) + 2)
        dc = DIST_SQUARE * (max(cols) - min(cols) + 2)
        dl = 0.34 * (max(lens) + 2)

        return 2 * max([dr, dc, dl]) * BOX_FACTOR

    def __str__(self):
        a = '{\n"vstrands":[\n'
        if len(self.vhelices) > 0:
            for h in self.vhelices:
                a = a + str(h) + ','
            a = a[0:len(a) - 1]
        a = a + '}\n'
        return a

class vhelix (object):

    def __init__(self):
        self.stapLoop = []
        self.scafLoop = []
        self.skip = []
        self.loop = []
        self.stap_colors = []
        self.row = 0
        self.col = 0
        self.num = 0
        self.stap = []
        self.scaf = []
        self.cad_index = -1
        self.skiploop_bases = 0

    def get_length(self):
        return max (len(self.scaf), len(self.stap))

    len = property (get_length)

    def add_square(self, toadd, which):
        if which == 'stap':
            self.stap.append(toadd)
        elif which == 'scaf':
            self.scaf.append (toadd)
        else:
            base.Logger.log("Cannot add square that is not scaf or stap. Dying now", base.Logger.CRITICAL)
            sys.exit(-1)

    def __str__(self):
        a = '{\n'

        a = a + '"stapLoop":['
        if len(self.stapLoop) > 0:
            for i in self.stapLoop:
                a = a + str(i) + ','
            a = a[0:len(a) - 1]  # remove last comma
        a = a + '],\n'

        a = a + '"skip":['
        if len(self.skip) > 0:
            for e in self.skip:
                a = a + str(e) + ','
            a = a[0:len(a) - 1]  # remove last comma
        a = a + '],\n'

        a = a + '"loop":['
        if len(self.loop) > 0:
            for e in self.loop:
                a = a + str(e) + ','
            a = a[0:len(a) - 1]  # remove last comma
        a = a + '],\n'

        a = a + '"stap_colors":['
        if len (self.stap_colors) > 0:
            for e in self.stap_colors:
                a = a + str(e) + ','
            a = a[0:len(a) - 1]  # remove last comma
        a = a + '],\n'

        a = a + '"row":' + str(self.row) + ',\n'
        a = a + '"col":' + str(self.col) + ',\n'
        a = a + '"num":' + str(self.num) + ',\n'

        a = a + '"scafLoop":['
        if len(self.scafLoop) > 0:
            for i in self.scafLoop:
                a = a + str(i) + ','
            a = a[0:len(a) - 1]  # remove last comma
        a = a + '],\n'

        a = a + '"stap":['
        if len(self.stap) > 0:
            for i in self.stap:
                a = a + str(i) + ','
            a = a[0:len(a) - 1]  # remove last comma
        a = a + '],\n'

        a = a + '"scaf":['
        if len(self.scaf) > 0:
            for i in self.scaf:
                a = a + str(i) + ','
            a = a[0:len(a) - 1]  # remove last comma
        a = a + ']\n}'
        return a

class square(object):

    def __init__ (self, V_0=-1, b_0=-1, V_1=-1, b_1=-1):
        """
        V_0, b_0, V_1, b_1 are integer indices correspond to:
        virtual_helix_behind, virtual_base_behind, virtual_helix_ahead, virtual_base_ahead
        """
        self.V_0 = V_0
        self.b_0 = b_0
        self.V_1 = V_1
        self.b_1 = b_1

    def __str__ (self):
        return '[%i,%i,%i,%i]' % (self.V_0, self.b_0, self.V_1, self.b_1)

    def type(self, vhelix, myid):
        # find type of strand (junction) on this square
        # currently direction always equals zero...
        direction = 0
        if self.V_0 == -1 and self.b_0 == -1:
            if self.V_1 == -1 and self.b_1 == -1:
                return 'empty'
            elif self.V_1 == vhelix.num and abs(self.b_1 - myid) == 1:
                if direction == 0:
                    return 'begin'
                else:
                    return 'end'
        elif self.V_0 == vhelix.num and abs(self.b_0 - myid) == 1:
            if self.V_1 == -1:
                if direction == 0:
                    return 'end'
                else:
                    return 'begin'
            elif self.V_1 == vhelix.num and abs(self.b_1 - myid) == 1:
                return 'continue'
            else:
                # join
                if direction == 0:
                    return 'end'
                else:
                    return 'begin'
        else:
            if self.V_1 == vhelix.num and abs(self.b_1 - myid) == 1:
                if direction == 0:
                    return 'begin'
                else:
                    return 'end'

        # shouldn't get to here
        base.Logger.log('unexpected square array', base.Logger.WARNING)

def parse_cadnano(path):
    import json

    cadsys = vstrands()

    try:
        with open(path) as json_data:
            cadnano = json.load(json_data)
            for vstrand in cadnano["vstrands"]:
                vh = vhelix()
                for key, val in vstrand.items():
                    if key == "skip":
                        vh.skip = [abs(int(x)) for x in val]
                    else:
                        setattr(vh, key, val)
                vh.stap = [square(*i) for i in vh.stap]
                vh.scaf = [square(*i) for i in vh.scaf]
                vh.skiploop_bases = len(vh.skip) + sum(vh.loop) - sum(vh.skip)
                cadsys.add_vhelix(vh)
    except IOError:
        print >> sys.stderr, "File '" + path + "' not found, aborting"
        sys.exit(1)
    except ValueError:
        print >> sys.stderr, "Invalid json file '" + path + "', aborting"
        sys.exit(1)
    except:
        print >> sys.stderr, "Caught an error while parsing '" + path + "', aborting"
        sys.exit(1)

    return cadsys

if __name__ == '__main__':

    def print_usage():
        print >> sys.stderr, "USAGE:"
        print >> sys.stderr, "\t%s cadnano_file lattice_type" % sys.argv[0]
        print >> sys.stderr, "\t[-q\--sequence FILE] [-b\--box VALUE] [-e\--seed VALUE] [-p\--print-virt2nuc]"
        exit(1)

    if len(sys.argv) < 3:
        print_usage()

    shortArgs = 'q:b:e:p'
    longArgs = ['sequence=', 'box=', 'seed=', 'print-virt2nuc']

    side = False
    sequence_filename = False
    print_virt2nuc = False
    source_file = sys.argv[1]

    origami_sq = False
    origami_he = False
    if sys.argv[2] == "sq":
        origami_sq = True
    elif sys.argv[2] == "he":
        origami_he = True
    else:
        print >> sys.stderr, "Lattice_type should be either 'sq' or 'he'"
        exit(1)

    try:
        import getopt
        args, files = getopt.gnu_getopt(sys.argv[3:], shortArgs, longArgs)
        for k in args:
            if k[0] == '-q' or k[0] == "--sequence":
                sequence_filename = k[1]
            elif k[0] == '-b' or k[0] == "--box":
                side = float(k[1])
                base.Logger.log("The system will be put in a box of side %s (in oxDNA simulation units)" % str(side), base.Logger.INFO)
            elif k[0] == '-e' or k[0] == "--seed":
                np.random.seed(int(k[1]))
            elif k[0] == '-p' or k[0] == "--print-virt2nuc":
                print_virt2nuc = True

    except Exception:
        print_usage()

    vh_vb2nuc = cu.vhelix_vbase_to_nucleotide()
    vh_vb2nuc_final = cu.vhelix_vbase_to_nucleotide()

    cadsys = parse_cadnano(source_file)
    base.Logger.log("Using json file %s" % source_file, base.Logger.INFO)

    # define sequences by vhelix
    sequence_file = 0
    single_strand_system = False
    sequences = []
    block_seq = True
    if sequence_filename:
        sequence_file = open(sequence_filename, "r")
        base.Logger.log("Using sequence file '%s'" % sequence_filename, base.Logger.INFO)
        # with this we can remove all whitespace and we don't have issues with the different newline sequences (\n vs \r\n)
        pattern = re.compile('\s+')
        lines = sequence_file.readlines()
        for line in lines:
            seq = []
            for x in re.sub(pattern, '', line):
                if x in ["R", "r"]:
                    seq.append(np.random.randint(0, 4))
                else:
                    try:
                        seq.append(base.base_to_number[x])
                    except KeyError:
                        base.Logger.log("KeyError while converting base names to integer; check the sequence file", base.Logger.CRITICAL)
                        sys.exit()
            sequences.append(seq)
    else:
        base.Logger.log("No sequence file given, using random sequence", base.Logger.INFO)
        for ii in range(len(cadsys.vhelices)):
            seq = []
            for _ in range(cadsys.vhelices[ii].skiploop_bases):
                seq.append(np.random.randint(0, 4))
            sequences.append(seq)

    # check whether we're dealing with a 1 strand system (i.e. NOT double helix) across many vhelices and defined with 1 .sqs line
    if len(sequences) == 1 and len(cadsys.vhelices) > 1:
        base.Logger.log("One line detected in the sequence file. Since the cadnano file contains more than 1 virtual helix, the sequence found will be used as we were dealing with a single-strand system", base.Logger.INFO)
        single_strand_system = True
        block_seq = False

    vhelix_counter = 0
    if not side:
        side = cadsys.bbox()
        base.Logger.log("Using default box size, a factor %s larger than the size of the cadnano system" % str(BOX_FACTOR), base.Logger.INFO)
    vhelix_direction_initial = np.array([0., 0., 1.])
    vhelix_perp_initial = np.array([1., 0., 0.])
    if origami_sq:
        vhelix_perp_initial = vhelix_rotation_origami_sq(vhelix_direction_initial, vhelix_perp_initial)
    elif origami_he:
        vhelix_perp_initial = vhelix_rotation_origami_he(vhelix_direction_initial, vhelix_perp_initial)

    slice_sys = base.System([side, side, side])
    final_sys = base.System([side, side, side])
    strand_number = -1
    partner_list_scaf = []
    partner_list_stap = []
    found_partner = False
    join_list_scaf = []
    join_list_stap = []
    begin_helix = -1
    end_helix = -1
    for h in cadsys.vhelices:
        h.cad_index = vhelix_counter
        if origami_sq:
            strands, helix_angles, pos, rot, vhelix_direction, vhelix_perp = generate_vhelices_origami_sq(vhelix_direction_initial, vhelix_perp_initial, h, sequence_file, single_strand_system, vhelix_counter)
        elif origami_he:
            strands, helix_angles, pos, rot, vhelix_direction, vhelix_perp = generate_vhelices_origami_he(vhelix_direction_initial, vhelix_perp_initial, h, sequence_file, single_strand_system, vhelix_counter)

        nodes = build_nodes(h)

        # read the scaffold squares and add strands to slice_sys
        i = 0
        for s in h.scaf:
            if s.V_0 == -1 and s.b_0 == -1:
                if s.V_1 == -1 and s.b_0 == -1:
                    pass
                elif s.V_1 == h.num and abs(s.b_1 - i) == 1:
                    if h.num % 2 == 0:
                        strand_number += 1
                    begin_helix = i
                    if h.num % 2 == 1:
                        slice_sys = add_slice(slice_sys, h, begin_helix, end_helix, nodes, strands, pos, vhelix_direction, vhelix_perp, rot, helix_angles, 0, block_seq, sequences)
                        vh_vb2nuc = add_slice_nupack(h, strand_number, begin_helix, end_helix, vh_vb2nuc, 2)
                else:
                    base.Logger.log("unexpected square array", base.Logger.WARNING)
            elif s.V_0 == h.num and abs(s.b_0 - i) == 1:
                if s.V_1 == -1 and s.b_1 == -1:
                    if h.num % 2 == 1:
                        strand_number += 1
                    end_helix = i
                    if h.num % 2 == 0:
                        slice_sys = add_slice(slice_sys, h, begin_helix, end_helix, nodes, strands, pos, vhelix_direction, vhelix_perp, rot, helix_angles, 0, block_seq, sequences)
                        vh_vb2nuc = add_slice_nupack(h, strand_number, begin_helix, end_helix, vh_vb2nuc, 2)
                elif s.V_1 == h.num and abs(s.b_1 - i) == 1:
                    pass
                else:
                    if h.num % 2 == 1:
                        strand_number += 1
                    end_helix = i
                    if h.num % 2 == 0 :
                        slice_sys = add_slice(slice_sys, h, begin_helix, end_helix, nodes, strands, pos, vhelix_direction, vhelix_perp, rot, helix_angles, 0, block_seq, sequences)
                        vh_vb2nuc = add_slice_nupack(h, strand_number, begin_helix, end_helix, vh_vb2nuc, 2)

                    if h.num % 2 == 1:
                        column = i
                    else:
                        column = i
                    for j in range(len(partner_list_scaf)):

                        if [h.num, column] == partner_list_scaf[j]:
                            join_list_scaf[j].insert(0, strand_number)
                            found_partner = True
                    if found_partner == False:
                        join_list_scaf.append([strand_number])
                        partner_list_scaf.append([s.V_1, s.b_1])
                    found_partner = False
            else:
                if s.V_1 == -1 and s.b_1 == -1:
                    base.Logger.log("unexpected square array", base.Logger.WARNING)
                elif s.V_1 == h.num and abs(s.b_1 - i) == 1:
                    if h.num % 2 == 0:
                        strand_number += 1
                    begin_helix = i
                    if h.num % 2 == 1:
                        slice_sys = add_slice(slice_sys, h, begin_helix, end_helix, nodes, strands, pos, vhelix_direction, vhelix_perp, rot, helix_angles, 0, block_seq, sequences)
                        vh_vb2nuc = add_slice_nupack(h, strand_number, begin_helix, end_helix, vh_vb2nuc, 2)

                    for j in range(len(partner_list_scaf)):
                        if h.num % 2 == 1:
                            column = i
                        else:
                            column = i
                        if [h.num, column] == partner_list_scaf[j]:
                            join_list_scaf[j].append(strand_number)
                            found_partner = True
                    if found_partner == False:
                        join_list_scaf.append([strand_number])
                        partner_list_scaf.append([s.V_0, s.b_0])
                    found_partner = False
                else:
                    base.Logger.log("unexpected square array", base.Logger.WARNING)
            i += 1

        if slice_sys.N_strands == 0:
            base.Logger.log("No scaffold strand found in virtual helix n. %d: staples-only virtual helices are not supported" % h.num, base.Logger.WARNING)
            continue

        # read the staple squares and add strands to slice_sys
        i = 0
        for s in h.stap:
            if s.V_0 == -1 and s.b_0 == -1:
                if s.V_1 == -1 and s.b_0 == -1:
                    pass
                elif s.V_1 == h.num and abs(s.b_1 - i) == 1:
                    if h.num % 2 == 1:
                        strand_number += 1
                    begin_helix = i
                    if h.num % 2 == 0:
                        slice_sys = add_slice(slice_sys, h, begin_helix, end_helix, nodes, strands, pos, vhelix_direction, vhelix_perp, rot, helix_angles, 1, block_seq, sequences)
                        vh_vb2nuc = add_slice_nupack(h, strand_number, begin_helix, end_helix, vh_vb2nuc, 3)
                else:
                    base.Logger.log("unexpected square array", base.Logger.WARNING)
            elif s.V_0 == h.num and abs(s.b_0 - i) == 1:
                if s.V_1 == -1 and s.b_1 == -1:
                    if h.num % 2 == 0:
                        strand_number += 1
                    end_helix = i
                    if h.num % 2 == 1:
                        slice_sys = add_slice(slice_sys, h, begin_helix, end_helix, nodes, strands, pos, vhelix_direction, vhelix_perp, rot, helix_angles, 1, block_seq, sequences)
                        vh_vb2nuc = add_slice_nupack(h, strand_number, begin_helix, end_helix, vh_vb2nuc, 3)

                elif s.V_1 == h.num and abs(s.b_1 - i) == 1:
                    pass
                else:
                    if h.num % 2 == 0:
                        strand_number += 1
                    end_helix = i
                    if h.num % 2 == 1:
                        slice_sys = add_slice(slice_sys, h, begin_helix, end_helix, nodes, strands, pos, vhelix_direction, vhelix_perp, rot, helix_angles, 1, block_seq, sequences)
                        vh_vb2nuc = add_slice_nupack(h, strand_number, begin_helix, end_helix, vh_vb2nuc, 3)

                    if h.num % 2 == 0:
                        column = i
                    else:
                        column = i
                    for j in range(len(partner_list_stap)):

                        if [h.num, column] == partner_list_stap[j]:
                            join_list_stap[j].insert(0, strand_number)
                            found_partner = True
                    if found_partner == False:
                        join_list_stap.append([strand_number])
                        partner_list_stap.append([s.V_1, s.b_1])
                    found_partner = False
            else:
                if s.V_1 == -1 and s.b_1 == -1:
                    base.Logger.log("unexpected square array", base.Logger.WARNING)
                elif s.V_1 == h.num and abs(s.b_1 - i) == 1:
                    if h.num % 2 == 1:
                        strand_number += 1
                    begin_helix = i
                    if h.num % 2 == 0:
                        slice_sys = add_slice(slice_sys, h, begin_helix, end_helix, nodes, strands, pos, vhelix_direction, vhelix_perp, rot, helix_angles, 1, block_seq, sequences)
                        vh_vb2nuc = add_slice_nupack(h, strand_number, begin_helix, end_helix, vh_vb2nuc, 3)

                    for j in range(len(partner_list_stap)):
                        if h.num % 2 == 0:
                            column = i
                        else:
                            column = i
                        if [h.num, column] == partner_list_stap[j]:
                            join_list_stap[j].append(strand_number)
                            found_partner = True
                    if found_partner == False:
                        join_list_stap.append([strand_number])
                        partner_list_stap.append([s.V_0, s.b_0])
                    found_partner = False
                else:
                    base.Logger.log("unexpected square array", base.Logger.WARNING)
            i += 1
        vhelix_counter += 1

    join_lists = [join_list_scaf, join_list_stap]

    # add strands to final_sys that aren't joined
    join_list_unpacked = []
    for a in range(2):
        for i in join_lists[a]:
            join_list_unpacked.extend(i)
    for i in range(len(slice_sys._strands)):
        if i not in join_list_unpacked:
            final_sys.add_strand(slice_sys._strands[i], check_overlap=False)
            vh_vb2nuc_final.add_strand(i, vh_vb2nuc)

    for a in range(2):
        join_list = join_lists[a]
        all_are_joined = False
        restart = False

        # check distance between the backbones we are about to join
        for pair in join_list:
            strand1 = slice_sys._strands[pair[0]]
            strand2 = slice_sys._strands[pair[1]]
            backbone_backbone_dist = strand1._nucleotides[-1].distance(strand2._nucleotides[0], PBC=False)
            absolute_bb_dist = np.sqrt(np.dot(backbone_backbone_dist, backbone_backbone_dist))
            if absolute_bb_dist > 1.0018 or absolute_bb_dist < 0.5525:
                base.Logger.log("the backbone-backbone distance across joints is %f: it will have to be relaxed with preliminary simulations" % absolute_bb_dist, base.Logger.WARNING)

        # match up all the pairs of joins that involve the same strand
        circular = []
        while all_are_joined == False:
            restart = False
            for i in range(len(join_list)):
                if restart == True:
                    break
                for j in range(len(join_list)):
                    if restart == True:
                        break
                    if join_list[i][0] == join_list[j][-1]:
                        if i != j:
                            join_list[j].extend(join_list[i][1:])
                            join_list.pop(i)
                            restart = True
                            break
                        else:
                            if i not in circular:
                                circular.append(i)

            if restart == False:
                all_are_joined = True

        # add joined strands
        for ii, join in enumerate(join_list):
            joined_strand = slice_sys._strands[join[0]]
            if ii in circular:
                for k in range(1, len(join) - 1):
                    joined_strand = joined_strand.append(slice_sys._strands[join[k]])
                joined_strand.make_circular(check_join_len=True)
            else:
                for k in range(1, len(join)):
                    joined_strand = joined_strand.append(slice_sys._strands[join[k]])

            final_sys.add_strand(joined_strand, check_overlap=False)

            # This is a bug fix. Ben 12/2/14
            # for a circular strand we need to terminate the strand one element early (so reduce the length
            # of the range by 1), since the final element is just a repeat of the first one.
            if joined_strand._circular:
                joining_range = range(len(join) - 2)
            else:
                joining_range = range(len(join) - 1)
            # add joined strands to v2n index
            for k in joining_range:
                vh_vb2nuc_final.add_strand(join[k], vh_vb2nuc, continue_join=True)
            vh_vb2nuc_final.add_strand(join[k + 1], vh_vb2nuc, continue_join=False)

            if single_strand_system == 1:
                final_sys._strands[0].set_sequence(sequences[0])

    if sequence_file and single_strand_system:
        if len(final_sys._strands) > 1:
            base.Logger.log("more than one strand detected - sequence file will not be read", base.Logger.WARNING)
            final_sys._strands[0].set_sequence(np.random.randint(0, 4, len(final_sys._strands[0]._nucleotides)))  # this line does not work

    # # Fix to reverse the direction of every strand so that the 3' to 5' direction is the same
    # # as in Cadnano. In cadnano the strands point in the 5' to 3' direction, whereas in oxDNA
    # # they point in the 3' to 5' direction. Ben 29/11/13
    rev_sys = base.System(final_sys._box)
    for strand in final_sys._strands:
        reverse_nucs = [nuc for nuc in strand._nucleotides]
        reverse_nucs.reverse()
        rev_strand = base.Strand()
        for nuc in reverse_nucs:
            rev_strand.add_nucleotide(base.Nucleotide(nuc.cm_pos, nuc._a1, -nuc._a3, nuc._base, nuc._btype))
        if strand._circular:
            rev_strand.make_circular(check_join_len=True)
        rev_sys.add_strand(rev_strand, check_overlap=False)
    # # also reverse the vhelix_vbase_to_nucleotide order so it corresponds to the reversed system
    vh_vb2nuc_rev = cu.vhelix_vbase_to_nucleotide()
    # count the number of nucleotides up to but not including the nucleotides in strand ii
    nnucs_to_here = range(rev_sys._N_strands)
    nuc_total = 0
    for strandii, strand in enumerate(rev_sys._strands):
        nnucs_to_here[strandii] = nuc_total
        nuc_total += len(strand._nucleotides)

    # fill in the _scaf and _stap dicts for the reverse vhelix_vbase_to_nucleotide object
    for vh, vb in vh_vb2nuc_final._scaf.keys():
        strandii, nuciis = vh_vb2nuc_final._scaf[(vh, vb)]
        rev_nuciis = []
        for nucii in nuciis:
            rev_nuciis.append(len(rev_sys._strands[strandii]._nucleotides) - 1 - (nucii - nnucs_to_here[strandii]) + nnucs_to_here[strandii])
        vh_vb2nuc_rev.add_scaf(vh, vb, strandii, rev_nuciis)
    for vh, vb in vh_vb2nuc_final._stap.keys():
        strandii, nuciis = vh_vb2nuc_final._stap[(vh, vb)]
        rev_nuciis = []
        for nucii in nuciis:
            rev_nuciis.append(len(rev_sys._strands[strandii]._nucleotides) - 1 - (nucii - nnucs_to_here[strandii]) + nnucs_to_here[strandii])
        vh_vb2nuc_rev.add_stap(vh, vb, strandii, rev_nuciis)

    # dump the spatial arrangement of the vhelices to a file
    vhelix_pattern = {}
    for i in range(len(cadsys.vhelices)):
        vhelix_pattern[cadsys.vhelices[i].num] = (cadsys.vhelices[i].row,cadsys.vhelices[i].col)

    if rev_sys.N == 0:
        base.Logger.log("The generated configuration is empty: this might be due to this conversion module not supporting virtual helices containing no scaffold strands.", base.Logger.CRITICAL)
        exit(1)

    if print_virt2nuc:
        with open("virt2nuc", "w") as fout:
            pickle.dump((vh_vb2nuc_rev, vhelix_pattern), fout)
            print >> sys.stderr, "## Wrote nucleotides' index conversion data to virt2nuc"

    basename = os.path.basename(sys.argv[1])
    topology_file = basename + ".top"
    configuration_file = basename + ".oxdna"

    rev_sys.print_lorenzo_output(configuration_file, topology_file)

    print >> sys.stderr, "## Wrote data to '%s' / '%s'" % (configuration_file, topology_file)
    print >> sys.stderr, "## DONE"