Converts to Rust and adds modern features

#!/usr/bin/env python

"""
Small projection test to compare with IDL tomograpy.
"""

import numpy as np
import tomograpy
im = tomograpy.centered_stack(0.0016, 32, n_images=1, radius=200., fill=0.)
cube = tomograpy.centered_cubic_map(3, 256, fill=1.)
P = tomograpy.lo(im.header, cube.header, obstacle="sun")
im[:] = (P * cube.ravel()).reshape(im.shape)

#!/usr/bin/env python
import numpy as np
import os
import copy
import time
import tomograpy
import fitsarray as fa
import lo
import scipy.sparse.linalg as spl
# data
path = os.path.join(os.getenv('HOME'), 'data', '171dec08')
obsrvtry = 'STEREO_A'
time_window = ['2008-12-01T00:00:00.000', '2008-12-03T00:00:00.000']
# one image every time_step seconds
time_step = 4 * 3600.
data = tomograpy.secchi.read_data(path, bin_factor=4,
                               obsrvtry=obsrvtry,
                               time_window=time_window,
                               time_step=time_step)
# cube
shape = 3 * (128,)
header = {'CRPIX1':64., 'CRPIX2':64., 'CRPIX3':64.,
          'CDELT1':0.0234375, 'CDELT2':0.0234375, 'CDELT3':0.0234375,
          'CRVAL1':0., 'CRVAL2':0., 'CRVAL3':0.,}
cube = fa.zeros(shape, header=header)
# model
P = tomograpy.lo(data.header, cube.header)
D = [lo.diff(cube.shape, axis=i) for i in xrange(cube.ndim)]
hypers = cube.ndim * (1e0, )
# inversion
t = time.time()
A = P.T * P + np.sum([h * d.T * d for h, d in zip(hypers, D)])
b = P.T * data.flatten()
#callback = lo.iterative.CallbackFactory(verbose=True)
#x, info = spl.bicgstab(A, b, maxiter=100, callback=callback)
x, info = lo.acg(P, data.flatten(), D, hypers, maxiter=100,)
sol = cube.copy()
sol[:] = x.reshape(cube.shape)
print(time.time() - t)

#!/usr/bin/env python
import numpy as np
import os
import copy
import time
import tomograpy
import fitsarray as fa
import lo

# data
path = os.path.join(os.getenv('HOME'), 'data', 'tomograpy., 'cor1')
#obsrvtry = 'SOHO    '
#instrume = 'LASCO   '
time_window = ['2009/09/01 00:00:00.000', '2009/09/15 00:00:00.000']
time_step = 8 * 3600. # one image every time_step seconds
data = tomograpy.solar.read_data(path, bin_factor=8,
                               #time_window=time_window,
                               #time_step=time_step
                               )
# errors in data ...
data.header['BITPIX'][:] = -64
data[np.isnan(data)] = 0.
data.header['RSUN'] /= 16.
# cube
shape = np.asarray(3 * (128. ,))
crpix = shape / 2.
cdelt = 6. / shape
crval = np.zeros(3)
header = {'CRPIX1':crpix[0], 'CRPIX2':crpix[1], 'CRPIX3':crpix[2],
          'CDELT1':cdelt[0], 'CDELT2':cdelt[1], 'CDELT3':cdelt[2],
          'CRVAL1':0., 'CRVAL2':0., 'CRVAL3':0.,}
cube = fa.zeros(shape, header=header)
t = time.time()
cube = tomograpy.backprojector(data, cube, obstacle="sun")
print("backprojection time : " + str(time.time() - t))

# inversion
t = time.time()
u = .5
kwargs={
    "obj_rmin":1.5,
    "obj_rmax":3.,
    "data_rmin":1.5,
    "data_rmax":2.5,
    "mask_negative":True
}
P, D, obj_mask, data_mask = tomograpy.models.thomson(data, cube, u, **kwargs)
# bpj
b = data.flatten()
bpj = (P.T * b).reshape(cube.shape)
hypers = 1e3 * np.ones(3)
sol = lo.acg(P, b, D, hypers, maxiter=100, tol=1e-6)
print("inversion time : %f" % (time.time() - t))
# reshape solution
sol.resize(cube.shape)
sol = fa.asfitsarray(sol, header=cube.header)
# reproject solution
reproj = P * sol.ravel()
reproj.resize(data.shape)

#!/bin/env python
import numpy as np
import os
import copy
import time
import tomograpy
import fitsarray as fa
# data
path = os.path.join(os.getenv('HOME'), 'data', '171dec08')
obsrvtry = 'STEREO_A'
time_window = ['2008-12-01T00:00:00.000', '2008-12-03T00:00:00.000']
time_step = 4 * 3600. # one image every time_step seconds
data = tomograpy.secchi.read_data(path, bin_factor=4,
                               obsrvtry=obsrvtry,
                               time_window=time_window,
                               time_step=time_step)
# cube
shape = 3 * (128,)
header = dict()
for i in xrange(1, 4):
    header['CRPIX' + str(i)] = 64.
    header['CDELT' + str(i)] = 0.0234375
    header['CRVAL' + str(i)] = 0.
cube = fa.zeros(shape, header=header, dtype=np.float32)
P = tomograpy.sun_lo(data.header, cube.header)
t = time.time()
fbp = (P.T * data.flatten()).reshape(cube.shape)
print("backprojection time : " + str(time.time() - t))

t = time.time()
fbp0 = tomograpy.backprojector_sun(data, cube)
print("backprojection time : " + str(time.time() - t))

#assert np.all(fbp == fbp0)

#!/usr/bin/env python
import numpy as np
import os
import copy
import time
import tomograpy
import fitsarray as fa
# data
path = os.path.join(os.getenv('HOME'), 'data', 'tomograpy., '171dec08')
obsrvtry = 'STEREO_A'
time_window = ['2008-12-01T00:00:00.000', '2008-12-15T00:00:00.000']
time_step = 8 * 3600. # one image every time_step seconds
data = tomograpy.solar.read_data(path, bin_factor=8,
                               obsrvtry=obsrvtry,
                               time_window=time_window,
                               time_step=time_step)
# map
cube = tomograpy.centered_cubic_map(3, 128, fill=0.)
t = time.time()
cube = tomograpy.backprojector(data, cube, obstacle="sun")
print("backprojection time : " + str(time.time() - t))

#!/usr/bin/env python
import os
import time
import numpy as np
import lo
import tomograpy
from tomograpy.solar import read_data
# data
obsrvtry = ('STEREO_A', 'STEREO_B')
data = tomograpy.solar.concatenate(
    [read_data(os.path.join(os.getenv('HOME'), 'data', 'tomograpy., '171dec08'),
               bin_factor=4,
               obsrvtry=obs,
               time_window=['2008-12-01T00:00:00.000',
                            '2008-12-15T00:00:00.000'],
               time_step= 4 * 3600.
               )
     for obs in obsrvtry])
data = tomograpy.solar.sort_data_array(data)
# scale A and B images
# the ratio of sensitivity between EUVI A and B
calibration_ba = {171:0.902023, 195:0.974536, 284:0.958269, 304:1.05954}
for i in xrange(data.shape[-1]):
    if data.header['OBSRVTRY'][i] == 'STEREO_B':
        data[..., i] /= calibration_ba[data.header['WAVELNTH'][i]]

# make sure it is 64 bits data
data.header['BITPIX'][:] = -64
# cube
shape = 3 * (128,)
header = {'CRPIX1':64.,
          'CRPIX2':64.,
          'CRPIX3':64.,
          'CDELT1':0.0234375,
          'CDELT2':0.0234375,
          'CDELT3':0.0234375,
          'CRVAL1':0.,
          'CRVAL2':0.,
          'CRVAL3':0.,}
cube = tomograpy.fa.zeros(shape, header=header)
# model
kwargs = {'obj_rmin':1., 'obj_rmax':1.4, 'data_rmax':1.3,
          'mask_negative':True, 'dt_min':100}
P, D, obj_mask, data_mask = tomograpy.models.stsrt(data, cube, **kwargs)
# apply mask to data
data *= (1 - data_mask)
# hyperparameters
hypers = (1e-1, 1e-1, 1e-1, 1e6)
# test time for one projection
t = time.time()
u = P.T * data.ravel()
print("maximal time : %f" % ((time.time() - t) * 100))
# inversion
t = time.time()
b = data.ravel()
#sol = lo.acg(P, b, D, hypers, maxiter=100)
sol = lo.rls(P, b, D, hypers, maxiter=100)
# reshape result
fsol = tomograpy.fa.asfitsarray(sol.reshape(obj_mask.shape), header=header)
print(time.time() - t)
fsol.tofits('stsrt_test.fits')

#!/usr/bin/env python
import os
import time
import numpy as np
import lo
import tomograpy
from tomograpy.solar import read_data
# data
data = read_data(os.path.join(os.getenv('HOME'), 'data', 'tomograpy., '171dec08'),
                 bin_factor=4.,
                 time_window=['2008-12-01T00:00:00.000',
                              '2008-12-15T00:00:00.000'],
                 time_step=8 * 3600.
                 )
data = tomograpy.solar.sort_data_array(data)
# scale A and B images
# the ratio of sensitivity between EUVI A and B
calibration_ba = {171:0.902023, 195:0.974536, 284:0.958269, 304:1.05954}
for i in xrange(data.shape[-1]):
    if data.header[i]['OBSRVTRY'] == 'STEREO_B':
        data[..., i] /= calibration_ba[data.header[i]['WAVELNTH']]

    # make sure it is 64 bits data
    data.header[i]['BITPIX'] = -64

# cube
cube = tomograpy.centered_cubic_map(3, 64, fill=0.)
# model
kwargs = {'obj_rmin':1., 'obj_rmax':1.5, 'data_rmin':1., 'data_rmax':1.3,
          'mask_negative':False, 'mask_nan':True}
P, D, obj_mask, data_mask = tomograpy.models.srt(data, cube, **kwargs)
# apply mask to data
data *= (1 - data_mask)
data[np.isnan(data)] = 0.
# hyperparameters
hypers = cube.ndim * (1e0, )
# inversion
# expected time
b = data.ravel()
b[np.isnan(b)] = 0.
t = time.time()
bpj = P.T * b
print((time.time() - t) * 4 * 100 )
# real time
t = time.time()
sol = lo.acg(P, b, D, hypers, maxiter=100)
print(time.time() - t)
# reshape result
fsol = tomograpy.fa.asfitsarray(sol.reshape(cube.shape), header=cube.header)
tomograpy.path = os.path.join(os.getenv('HOME'), 'data', 'tomograpy.)
fsol.tofits(os.path.join(tomograpy.path, "output", "test_tomograpy.secchi_mask.fits"))

#!/usr/bin/env python
import time
import numpy as np
import tomograpy
# object
object_header = tomograpy.centered_cubic_map_header(3, 128)
obj = tomograpy.simu.object_from_header(object_header, fill=1.)
# data
radius = 200.
a = tomograpy.fov(object_header, radius)
data = tomograpy.centered_stack(a, 128, n_images=60, radius=200., max_lon=np.pi)
# projection
t = time.time()
data = tomograpy.projector(data, obj)
print("projection time : " + str(time.time() - t))
# backprojection
t = time.time()
data[:] = 1.
obj0 = tomograpy.simu.object_from_header(object_header, fill=0.)
obj0 = tomograpy.backprojector(data, obj0)
print("backprojection time : " + str(time.time() - t))

obj1 = tomograpy.simu.object_from_header(object_header, fill=0.)
obj1 = tomograpy.backprojector(data, obj1)

#!/usr/bin/env python
import time
import numpy as np
import tomograpy
# object
obj = tomograpy.centered_cubic_map(3, 128, fill=1.)
# number of images
n = 20
# reshape object for 4d model
obj4 = obj.reshape(obj.shape + (1,)).repeat(n, axis=-1)
obj4.header['NAXIS'] = 4
obj4.header['NAXIS4'] = obj4.shape[3]
obj4.header['CRVAL4'] = 0.

# data
radius = 200
a = tomograpy.fov(obj.header, radius)
data = tomograpy.centered_stack(a, 128, n_images=n, radius=radius,
                                    max_lon=np.pi)
data[:] = np.zeros(data.shape)
# projection
t = time.time()
data = tomograpy.projector4d(data, obj4)
print("projection time : " + str(time.time() - t))
# backprojection
x0 = obj4.copy()
x0[:] = 0.
t = time.time()
x0 = tomograpy.backprojector4d(data, x0)
print("backprojection time : " + str(time.time() - t))

#!/usr/bin/env python
import time
import numpy as np
import tomograpy
import lo
# object
obj = tomograpy.centered_cubic_map(3, 64, fill=1.)
# number of images
n = 64
# reshape object for 4d model
obj4 = obj[..., np.newaxis].repeat(n, axis=-1)
obj4.header['NAXIS'] = 4
obj4.header['NAXIS4'] = obj4.shape[3]
obj4.header['CRVAL4'] = 0.

# data
radius = 200
a = tomograpy.fov(obj.header, radius)
data1 = tomograpy.centered_stack(a, 64, n_images=n/2, radius=radius,
                                    max_lon=np.pi, fill=0.)
data2 = tomograpy.centered_stack(a, 64, n_images=n/2, radius=radius,
                                    min_lon=np.pi / 2., max_lon=1.5 * np.pi, fill=0.)
data = tomograpy.solar.concatenate((data1, data2))

# times
DT = 1000.
dt_min = 100.
dates = np.arange(n / 2) * DT / 2.
dates = np.concatenate(2 * (dates, ))
dates = [time.strftime('%Y-%m-%dT%H:%M:%S', time.localtime((t))) for t in dates]
for i in xrange(len(data.header)):
    data.header[i]['DATE_OBS'] = dates[i]

data = tomograpy.solar.sort_data_array(data)

# projection
t = time.time()
data = tomograpy.projector4d(data, obj4, obstacle="sun")
print("projection time : " + str(time.time() - t))
# backprojection
x0 = obj4.copy()
x0[:] = 0.
t = time.time()
x0 = tomograpy.backprojector4d(data, x0, obstacle="sun")
print("backprojection time : " + str(time.time() - t))

# model
kwargs = {'obj_rmin':1., #'obj_rmax':1.3,
          'mask_negative':False, 'dt_min':100}
P, D, obj_mask, data_mask = tomograpy.models.stsrt(data, obj, **kwargs)
# hyperparameters
hypers = (1e-1, 1e-1, 1e-1, 1e3)
# test time for one projection
b = data.ravel()
t = time.time()
u = P.T * b
print("time with index grouping : %f" % ((time.time() - t)))
# inversion
t = time.time()
sol = lo.acg(P, b, D, hypers, tol=1e-10, maxiter=100)
# reshape result
fsol = tomograpy.fa.asfitsarray(sol.reshape(obj_mask.shape), header=obj4.header)
print(time.time() - t)
#fsol.tofits('stsrt_test.fits')

#!/usr/bin/env python
import time
import numpy as np
import tomograpy
import lo
# object
obj = tomograpy.centered_cubic_map(3, 32)
obj[:] = tomograpy.phantom.shepp_logan(obj.shape)
# data
radius = 200.
a = tomograpy.fov(obj.header, radius)
data = tomograpy.centered_stack(a, 128, n_images=60, radius=radius,
                                    max_lon=np.pi)
# projector
P = tomograpy.lo(data.header, obj.header)
# projection
t = time.time()
data = tomograpy.projector(data, obj)
print("projection time : " + str(time.time() - t))
# data
y = data.flatten()
# backprojection
t = time.time()
x0 = P.T * y
bpj = x0.reshape(obj.shape)
print("projection time : " + str(time.time() - t))
# priors
Ds = [lo.diff(obj.shape, axis=i) for i in xrange(3)]
# inversion using scipy.sparse.linalg
t = time.time()
sol = lo.acg(P, y, Ds, 1e-2 * np.ones(3),  maxiter=100, tol=1e-20)
sol = sol.reshape(bpj.shape)
print("inversion time : " + str(time.time() - t))

#!/usr/bin/env python
import time
import numpy as np
import tomograpy
# object
obj = tomograpy.centered_cubic_map(3, 128, fill=1.)
# data
radius = 200.
a = tomograpy.fov(obj.header, radius)
data = tomograpy.centered_stack(a, 128, n_images=17, radius=200., max_lon=np.pi)
# projection
t = time.time()
data = tomograpy.projector(data, obj, obstacle="sun")
print("projection time : " + str(time.time() - t))
# backprojection
obj0 = tomograpy.centered_cubic_map(3, 128, fill=0.)
t = time.time()
obj0 = tomograpy.backprojector(data, obj0, obstacle="sun")
print("backprojection time : " + str(time.time() - t))

#!/usr/bin/env python
import time
import numpy as np
import tomograpy
import lo
# object
obj = tomograpy.centered_cubic_map(10, 64)
obj[:] = tomograpy.phantom.shepp_logan(obj.shape)
# data
radius = 200
a = tomograpy.fov(obj, radius)
data = tomograpy.centered_stack(a, 128, n_images=60, radius=radius, max_lon=np.pi)
# model
kwargs = {"pb":"pb", "obj_rmin":1.5, "data_rmin":1.5}
P, D, obj_mask, data_mask = tomograpy.models.thomson(data, obj, u=.5, **kwargs)
# projection
t = time.time()
data[:] = (P * obj.ravel()).reshape(data.shape)
print("projection time : " + str(time.time() - t))
# data
# backprojection
t = time.time()
x0 = P.T * data.ravel()
bpj = x0.reshape(obj.shape)
print("backprojection time : " + str(time.time() - t))
# inversion using scipy.sparse.linalg
t = time.time()
sol = lo.acg(P, data.ravel(), D, 1e-3 * np.ones(3),  maxiter=100, tol=1e-8)
sol = sol.reshape(obj.shape)
print("inversion time : " + str(time.time() - t))

import matplotlib.pyplot as plt
import numpy as np
import pylops

from tomograpy import project_3d, backproject_3d

if __name__ == "__main__":
    print("Test started")

    img_size = 50
    cube_size = 50
    x, y = np.meshgrid(np.arange(img_size, dtype=np.float32), np.arange(img_size, dtype=np.float32))
    z = np.zeros((img_size, img_size), dtype=np.float32) + 3
    densities = np.ones((cube_size, cube_size, cube_size), dtype=np.float32)
    # densities[1, 4, 1] = 100
    densities[2, 2, 2] = 5
    densities[8, 8, 2] = 8
    densities[10:30, 10:30, 0:30] = 3
    mask = np.ones((cube_size, cube_size, cube_size), dtype=bool)

    #b = (0.5, 0.5, 0.5)
    b = (0, 0, 0)
    #b = (1E-8, 1E-8, 1E-8)
    #b = (1, 1, 1)
    #b = (10, 10, 10)
    delta = (1, 1, 1)
    #unit_normal = (-1.0, -1E-8, -1E-8)
    #unit_normal = (1.0, 1E-8, 1E-8)
    #unit_normal = (-1, 0, 0)
    #unit_normal = (-1/np.sqrt(3), -1/np.sqrt(3), -1/np.sqrt(3))

    unit_normal = (1E-7, 1E-7, -1.0)
    path_distance = 500

    result = project_3d(x, y, z, densities, mask, b, delta, unit_normal, path_distance)
    # print(result)

    import pylops

    def forward(densities):
        return project_3d(x, y, z, densities.reshape((cube_size, cube_size, cube_size)).astype(np.float32),
                          mask, b, delta, unit_normal, path_distance).reshape((img_size**2))

    def backward(image):
        return backproject_3d(x, y, z, image.reshape((img_size, img_size)).astype(np.float32),
                   np.zeros_like(densities, dtype=np.float32),
                   mask, b, delta, unit_normal, path_distance, True).reshape((cube_size**3))

    def test_forward(xarr):
        print(xarr)
        return forward(xarr)

    from pylops.basicoperators import FunctionOperator

    op = pylops.FunctionOperator(forward, backward, img_size**2, cube_size**3)
    out = op @ densities.reshape((cube_size**3))

    fig, ax = plt.subplots()
    ax.imshow(out.reshape((img_size, img_size)))
    fig.show()
    # print(out.shape)
    # image = 2*np.ones((5, 5), dtype=np.float32)
    # image[2, 2] = 10

    image = out.reshape((img_size, img_size))

    import xarray as xr

    da = xr.DataArray(
        data=image.reshape((img_size*img_size)),
        dims = ["x"],
        coords = dict(
            lon=(["x"], x.reshape((img_size*img_size)))
        ),
        # dims=["x", "y"],
        # coords=dict(
        #     lon=(["x", "y"], x.reshape((img_size*img_size))),
        #     lat=(["x", "y"], y.reshape((img_size*img_size))),
        # ),
        attrs=dict(
            description="Ambient temperature.",
            units="degC",
        ),
    )

    xinv = op / image.reshape((img_size*img_size))

    # xinv = pylops.optimization.leastsquares.regularized_inversion(
    #     op, image.reshape((img_size*img_size)), [], **dict(damp=0, iter_lim=10, show=True, atol=1E-8, btol=1E-8)
    # )[0]
    # print(xinv.reshape((10, 10, 10)))

    fig, axs = plt.subplots(ncols=2, nrows=3, figsize=(10, 14))
    im = axs[0, 0].imshow(densities[:, :, 2], vmin=0, vmax=3)
    axs[0, 0].set_title("Input density slice")
    fig.colorbar(im)

    im = axs[1, 0].imshow(xinv.reshape((cube_size, cube_size, cube_size))[:, :, 2], vmin=0, vmax=3)
    axs[1, 0].set_title("Reconstructed density slice")
    fig.colorbar(im)

    im = axs[2, 0].imshow(densities[:, :, 2] - xinv.reshape((cube_size, cube_size, cube_size))[:, :, 2], vmin=-3, vmax=3, cmap='seismic')
    axs[2, 0].set_title("Input - reconstruction")
    fig.colorbar(im)

    im = axs[1, 1].imshow(image, vmin=0, vmax=100)
    axs[1, 1].set_title("Image used in reconstruction")
    fig.colorbar(im)

    axs[0, 1].set_axis_off()
    axs[2, 1].set_axis_off()
    plt.show()
    # BACKPROJECT
    # image = 2*np.ones((15, 15), dtype=np.float32)
    # result = backproject_3d(x, y, z, image,
    #                         np.zeros_like(densities, dtype=np.float32),
    #                         mask, b, delta, unit_normal, path_distance, True)
    # print("cube", result)
    #
    # import matplotlib as mpl
    # import matplotlib.pyplot as plt
    #
    # mpl.use('macosx')
    #
    # fig = plt.figure()
    # ax = fig.add_subplot(projection='3d')
    # x, y, z = np.where(result > 1)
    # ax.scatter(x, y, z)
    #
    # plt.show(block=True)

    print("Test ended")

import matplotlib.pyplot as plt
import numpy as np
import pylops
from datetime import datetime

from tomograpy import project_3d, backproject_3d
from pylops.basicoperators import FunctionOperator
from pylops import LinearOperator, lsqr

def z_rotation_matrix_3d(angle):
    return np.array([[[np.cos(angle), np.sin(angle), 0],
                      [-np.sin(angle), np.cos(angle), 0],
                      [0, 0, 1]]])

if __name__ == "__main__":
    print("Test started")
    start_time = datetime.now()

    angles = np.linspace(0, np.pi, 14)

    radius = 300

    img_size = 20
    cube_size = 20
    densities = np.zeros((cube_size, cube_size, cube_size), dtype=np.float32)
    densities[3:-3, 3:-3, 3:-3] = 100
    densities[7:-7, 7:-7, 7:-7] = 0
    # densities = np.random.randint(0, 100, (cube_size, cube_size, cube_size)).astype(np.float32)
    # densities[5, 5, :] = 0

    mask = np.ones((cube_size, cube_size, cube_size), dtype=bool)

    b = (0, 0, 0)
    b = (-cube_size / 2, -cube_size / 2, -cube_size / 2)

    delta = (1.0, 1.0, 1.0)
    path_distance = 500.0

    norms, xs, ys, zs, ds, imgs = [], [], [], [], [], []
    for angle in angles:
        t_angle = -angle + np.pi/2

        img_x = np.arange(img_size) - img_size / 2
        img_y = np.zeros((img_size, img_size))
        img_z = np.arange(img_size) - img_size / 2
        img_x, img_z = np.meshgrid(img_x, img_z)

        img_x, img_y, img_z = img_x.flatten(), img_y.flatten(), img_z.flatten()

        R = z_rotation_matrix_3d(t_angle)
        coords = (R @ np.stack([img_x, img_y, img_z]))[0]
        img_x, img_y, img_z = coords[0], coords[1], coords[2]
        img_x = radius * np.cos(angle) + img_x
        img_y = radius * np.sin(angle) + img_y

        xx = img_x.reshape((img_size, img_size)).astype(np.float32)
        yy = -img_y.reshape((img_size, img_size)).astype(np.float32)
        zz = img_z.reshape((img_size, img_size)).astype(np.float32)

        v1 = np.array([xx[0, 1] - xx[0, 0], yy[0, 1] - yy[0, 0], zz[0, 1] - zz[0, 0]])
        v2 = np.array([xx[1, 0] - xx[0, 0], yy[1, 0] - yy[0, 0], zz[1, 0] - zz[0, 0]])
        v1 = v1 / np.linalg.norm(v1)
        v2 = v2 / np.linalg.norm(v2)
        normal = np.cross(v1, v2)
        normal = normal / np.linalg.norm(normal)

        norm = normal
        norm[norm == 0] = 1E-6
        norms.append(norm)
        print(np.rad2deg(angle), np.rad2deg(t_angle), norm)

        d = 500
        img = project_3d(xx, yy, zz, densities, mask, b, delta, norm, d)
        xs.append(xx)
        ys.append(yy)
        zs.append(zz)
        ds.append(d)
        imgs.append(img.astype(np.float32))
    imgs = np.array(imgs)

    # show
    for angle, img in zip(angles, imgs):
        fig, ax = plt.subplots()
        im = ax.imshow(img)
        fig.colorbar(im)
        fig.savefig(f"/Users/jhughes/Desktop/projection_dots/{int(np.rad2deg(angle)):03d}.png")
        plt.close()

    class Tomo(LinearOperator):
        def __init__(self, xs, ys, zs, norms, ds, b, delta, model_shape, mask, dtype=None):
            self.xs = xs
            self.ys = ys
            self.zs = zs
            self.norms = norms
            self.ds = ds
            self.b = b
            self.delta = delta
            self.model_shape = model_shape
            self.mask = mask
            super().__init__(dtype=np.dtype(dtype),
                             dims=self.model_shape,
                             dimsd=(1, self.xs[0].shape[0], self.xs[0].shape[1]))
        #                     dimsd=(len(self.xs), self.xs[0].shape[0], self.xs[0].shape[1]))
        #
        # def _matvec(self, densities):
        #     return np.array([project_3d(x, y, z, densities.reshape(self.model_shape).astype(np.float32),
        #                                 self.mask, self.b, self.delta, norm, d)
        #             for x, y, z, norm, d in zip(self.xs, self.ys, self.zs, self.norms, self.ds)]).flatten()
        #
        # def _rmatvec(self, imgs):
        #     densitiesi = np.zeros(self.model_shape, dtype=np.float32)
        #     for i, img in enumerate(imgs.reshape(len(self.xs), self.xs[0].shape[0], self.xs[0].shape[1])):
        #         # densitiesi += backproject_3d(self.xs[i], self.ys[i], self.zs[i], img,
        #         #                densitiesi,
        #         #                self.mask, self.b, self.delta, self.norms[i], self.ds[i], True)
        #         densitiesi = backproject_3d(self.xs[i], self.ys[i], self.zs[i], img,
        #                                      densitiesi,
        #                                      self.mask, self.b, self.delta, self.norms[i], self.ds[i], True)
        #     # return ((densitiesi - np.sum(img)) / (len(self.xs)-1)).astype(np.float32)
        #     return densitiesi / (cube_size - 1) / (len(self.xs) / 2 + 1)#/ len(self.xs) #/ densitiesi.shape[0] / len(self.xs)
        #     #return densitiesi.flatten().astype(np.float32)

        def _matvec(self, densities):
            # return np.array([project_3d(x, y, z, densities.reshape(self.model_shape).astype(np.float32),
            #                             self.mask, self.b, self.delta, norm, d)
            #         for x, y, z, norm, d in zip(self.xs, self.ys, self.zs, self.norms, self.ds)]).flatten()
            i = 0
            return project_3d(self.xs[i], self.ys[i], self.zs[i], densities.reshape(self.model_shape).astype(np.float32), self.mask, self.b, self.delta, self.norms[i], self.ds[i]).flatten()

        def _rmatvec(self, imgs):
            densitiesi = np.zeros(self.model_shape, dtype=np.float32)
            for i, img in enumerate(imgs.reshape(1, self.xs[0].shape[0], self.xs[0].shape[1])):
                # densitiesi += backproject_3d(self.xs[i], self.ys[i], self.zs[i], img,
                #                densitiesi,
                #                self.mask, self.b, self.delta, self.norms[i], self.ds[i], True)
                # if i == 0:
                densitiesi += backproject_3d(self.xs[i], self.ys[i], self.zs[i], img,
                                             np.zeros_like(densitiesi).astype(np.float32),
                                             self.mask, self.b, self.delta, self.norms[i], self.ds[i], True)
            # return ((densitiesi - np.sum(img)) / (len(self.xs)-1)).astype(np.float32)
            #return (densitiesi / (cube_size - 1)).flatten() #/ densitiesi.shape[0] / len(self.xs)
            return densitiesi.flatten().astype(np.float32)

    print(len(xs))
    op = Tomo(xs, ys, zs, norms, ds, b, delta, densities.shape, mask, dtype=np.float32)

    proj = op @ densities.flatten()
    densities_bpj = op.H @ proj.flatten()
    proj2 = op @ densities_bpj.flatten()
    densities_bpj2 = op.H @ proj2.flatten()
    print(np.where(proj == -999))

    #print(np.conj(proj).T @ (op @ densities))

    print("maxs", np.max(proj), np.max(proj2))
    print("pcts", np.nanpercentile(densities, 95), np.nanpercentile(densities_bpj, 95),  np.nanpercentile(densities_bpj2, 95))
    print("close", np.allclose(densities_bpj, densities_bpj2))

    fig, axs = plt.subplots(ncols=2)
    im = axs[0].imshow(proj.reshape((img_size, img_size)), vmin=0, vmax=1300)
    fig.colorbar(im)
    im = axs[1].imshow(proj2.reshape((img_size, img_size))/ (cube_size-1))#, vmin=0, vmax=1300)
    fig.colorbar(im)
    fig.savefig("/Users/jhughes/Desktop/projection_dots/comparison.png")

    fig, axs = plt.subplots(ncols=3)
    im = axs[0].imshow(densities.reshape((cube_size, cube_size, cube_size))[0], vmin=0, vmax=100)
    fig.colorbar(im)
    im = axs[1].imshow(densities_bpj.reshape((cube_size, cube_size, cube_size))[0])# , vmin=0, vmax=100)
    fig.colorbar(im)
    im = axs[2].imshow(densities_bpj2.reshape((cube_size, cube_size, cube_size))[0])# , vmin=0, vmax=100)
    fig.colorbar(im)
    fig.savefig("/Users/jhughes/Desktop/projection_dots/comparison_dense.png")

    from pylops.utils import dottest
    _ = dottest(op, 400, 8000, atol=0.1, complexflag=0, verb=True)

    model = pylops.optimization.basic.lsqr(op, imgs[0].flatten(), niter=10, show=True)[0]
    #
    # model = model.reshape(densities.shape)
    #
    # limit = np.nanpercentile(model, 95)
    #
    # for i in range(cube_size):
    #     fig, axs = plt.subplots(ncols=2)
    #     im = axs[0].imshow(densities[:, :, i], vmin=0, vmax=150)#, vmin=0, vmax=5)
    #     fig.colorbar(im)
    #     im = axs[1].imshow(model[:, :, i], vmin=0, vmax=limit)
    #     fig.colorbar(im)
    #     fig.show()
    #     fig.savefig(f"/Users/jhughes/Desktop/projection_dots/test_{i:03d}.png")

import matplotlib.pyplot as plt
import numpy as np
import pylops
import scipy.fftpack as fp
from tomograpy import project_3d, backproject_3d

def z_rotation_matrix_3d(angle):
    return np.array([[[np.cos(angle), np.sin(angle), 0],
                      [-np.sin(angle), np.cos(angle), 0],
                      [0, 0, 1]]])

angles = np.linspace(0, 2*np.pi, 50)

radius = 500

img_size = 100
cube_size = 100
# x, y = np.meshgrid(np.arange(img_size, dtype=np.float32), np.arange(img_size, dtype=np.float32))
# z = np.zeros((img_size, img_size), dtype=np.float32)
densities = np.zeros((cube_size, cube_size, cube_size), dtype=np.float32)
# densities[1:30, 1:30, 1:30] = 25
# densities[80:90, 80:90, 80:90] = 5
densities[10:-10, 10:-10, 10:-10] = 100
densities[30:-30, 30:-30, 30:-30] = 0

# densities[2, 2, 2] = 5
# densities[8, 8, 2] = 8
# densities[10:30, 10:30, 0:30] = 3
mask = np.ones((cube_size, cube_size, cube_size), dtype=bool)

b = (0, 0, 0)
b = (-cube_size / 2, -cube_size / 2, -cube_size / 2)

delta = (1.0, 1.0, 1.0)
# unit_normal = (1E-7, 1E-7, 1.0)
path_distance = 500.0

total = np.zeros_like(densities)

for angle in angles:
    t_angle = -angle + np.pi / 2  # np.deg2rad(np.abs(90*np.cos(angle)))

    img_x = np.arange(img_size) - img_size / 2
    img_y = np.zeros((img_size, img_size))
    img_z = np.arange(img_size) - img_size / 2
    img_x, img_z = np.meshgrid(img_x, img_z)

    img_x, img_y, img_z = img_x.flatten(), img_y.flatten(), img_z.flatten()

    R = z_rotation_matrix_3d(t_angle)
    coords = (R @ np.stack([img_x, img_y, img_z]))[0]
    img_x, img_y, img_z = coords[0], coords[1], coords[2]
    img_x = radius * np.cos(angle) + img_x  # - img_size/2
    img_y = radius * np.sin(angle) + img_y  # - img_size/2

    xx = img_x.reshape((img_size, img_size)).astype(np.float32)
    yy = -img_y.reshape((img_size, img_size)).astype(np.float32)
    zz = img_z.reshape((img_size, img_size)).astype(np.float32)

    v1 = np.array([xx[0, 1] - xx[0, 0], yy[0, 1] - yy[0, 0], zz[0, 1] - zz[0, 0]])
    v2 = np.array([xx[1, 0] - xx[0, 0], yy[1, 0] - yy[0, 0], zz[1, 0] - zz[0, 0]])
    v1 = v1 / np.linalg.norm(v1)
    v2 = v2 / np.linalg.norm(v2)
    normal = np.cross(v1, v2)
    normal = normal / np.linalg.norm(normal)

    norm = normal  # R @ np.array([1E-7, 1E-7, -1.0])
    norm[norm == 0] = 1E-6
    print(np.rad2deg(angle), np.rad2deg(t_angle), norm)

    proj = project_3d(xx, yy, zz, densities, mask, b, delta, norm, path_distance)

    fig, ax = plt.subplots()
    im = ax.imshow(proj, origin='lower')
    fig.colorbar(im)
    fig.savefig(f"/Users/jhughes/Desktop/projection_bpj/proj_{angle:0.3f}.png")
    plt.close()

    # FILTERING
    # (w, h) = proj.shape
    # half_w, half_h = int(w / 2), int(h / 2)
    #
    # F1 = fp.fft2(proj.astype(float))
    # F2 = fp.fftshift(F1)
    #
    # # high pass filter
    # n = 10
    # F2[half_w - n:half_w + n + 1, half_h - n:half_h + n + 1] = 0
    # proj = fp.ifft2(fp.ifftshift(F2)).real

    result = backproject_3d(xx, yy, zz, proj.astype(np.float32),
                       np.zeros_like(densities, dtype=np.float32),
                       mask, b, delta, norm, path_distance, True)
    print(result.shape)

    scaling = densities.shape[0]
    result = result / scaling

    total += result
# ((result - np.sum(proj)) / (len(self.xs)-1)).astype(np.float32)
#result = result - np.sum(proj)

# # Scale result to make sure that fbp(A, A(x)) == x holds at least
# # to some approximation. In limited experiments, this is true for
# # this version of FBP up to 1%.
# # *Note*: For some reason, we do not have to scale with respect to
# # the pixel dimension that is orthogonal to the rotation axis (`u`
# # or horizontal pixel dimension). Hence, we only scale with the
# # other pixel dimension (`v` or vertical pixel dimension).
# vg, pg = A.astra_compat_vg, A.astra_compat_pg
#
# pixel_height = (pg.det_size[0] / pg.det_shape[0])
# voxel_volume = np.prod(np.array(vg.size / np.array(vg.shape)))
# scaling = (np.pi / pg.num_angles) * pixel_height / voxel_volume
#
# rec *= scaling
# pixel_height = proj.size / proj.shape[0]
# voxel_volume = np.prod(densities.size / np.array(densities.shape))
# scaling = (np.pi / 1) * pixel_height / voxel_volume

total /= len(angles)

for i in range(cube_size):
    fig, axs = plt.subplots(ncols=2)
    im = axs[0].imshow(densities[i, :, :], origin='lower', vmin=0, vmax=150)
    fig.colorbar(im)

    im = axs[1].imshow(total[i, :, :], origin='lower')#, vmin=0, vmax=150)
    fig.colorbar(im)
    fig.savefig(f"/Users/jhughes/Desktop/projection_bpj/{i:03d}.png")
    plt.close()

import matplotlib.pyplot as plt
import numpy as np
import pylops
from datetime import datetime

from tomograpy import project_3d, backproject_3d
from pylops.basicoperators import FunctionOperator
from pylops import LinearOperator, lsqr

def z_rotation_matrix_3d(angle):
    return np.array([[[np.cos(angle), np.sin(angle), 0],
                      [-np.sin(angle), np.cos(angle), 0],
                      [0, 0, 1]]])

if __name__ == "__main__":
    print("Test started")
    start_time = datetime.now()

    angles = np.linspace(0, np.pi, 14) + np.pi/60

    radius = 300

    img_size = 100
    cube_size = 100
    # x, y = np.meshgrid(np.arange(img_size, dtype=np.float32), np.arange(img_size, dtype=np.float32))
    # z = np.zeros((img_size, img_size), dtype=np.float32)
    densities = np.zeros((cube_size, cube_size, cube_size), dtype=np.float32)
    # densities[1:30, 1:30, 1:30] = 25
    # densities[80:90, 80:90, 80:90] = 5
    densities[30:-30, 30:-30, 30:-30] = 100
    densities[40:-40, 40:-40, 40:-40] = 0

    # densities[2, 2, 2] = 5
    # densities[8, 8, 2] = 8
    # densities[10:30, 10:30, 0:30] = 3
    mask = np.ones((cube_size, cube_size, cube_size), dtype=bool)

    b = (0, 0, 0)
    b = (-cube_size / 2, -cube_size / 2, -cube_size / 2)

    delta = (1.0, 1.0, 1.0)
    # delta = (0.5, 0.5, 0.5)
    # unit_normal = (1E-7, 1E-7, 1.0)
    path_distance = 500.0

    norms, xs, ys, zs, ds, imgs = [], [], [], [], [], []
    #angles = np.array([0, 30])
    for angle in angles:
        t_angle = -angle + np.pi/2 #np.deg2rad(np.abs(90*np.cos(angle)))

        img_x = np.arange(img_size) - img_size / 2
        img_y = np.zeros((img_size, img_size))
        img_z = np.arange(img_size) - img_size / 2
        img_x, img_z = np.meshgrid(img_x, img_z)

        img_x, img_y, img_z = img_x.flatten(), img_y.flatten(), img_z.flatten()

        R = z_rotation_matrix_3d(t_angle)
        coords = (R @ np.stack([img_x, img_y, img_z]))[0]
        img_x, img_y, img_z = coords[0], coords[1], coords[2]
        img_x = radius * np.cos(angle) + img_x  # - img_size/2
        img_y = radius * np.sin(angle) + img_y  # - img_size/2

        xx = img_x.reshape((img_size, img_size)).astype(np.float32)
        yy = -img_y.reshape((img_size, img_size)).astype(np.float32)
        zz = img_z.reshape((img_size, img_size)).astype(np.float32)

        v1 = np.array([xx[0, 1] - xx[0, 0], yy[0, 1] - yy[0, 0], zz[0, 1] - zz[0, 0]])
        v2 = np.array([xx[1, 0] - xx[0, 0], yy[1, 0] - yy[0, 0], zz[1, 0] - zz[0, 0]])
        v1 = v1 / np.linalg.norm(v1)
        v2 = v2 / np.linalg.norm(v2)
        normal = np.cross(v1, v2)
        normal = normal / np.linalg.norm(normal)

        norm = normal  # R @ np.array([1E-7, 1E-7, -1.0])
        norm[norm == 0] = 1E-6
        norms.append(norm)
        print(np.rad2deg(angle), np.rad2deg(t_angle), norm)

        d = 500
        img = project_3d(xx, yy, zz, densities, mask, b, delta, norm, d)
        xs.append(xx)
        ys.append(yy)
        zs.append(zz)
        ds.append(d)
        imgs.append(img.astype(np.float32))
    imgs = np.array(imgs)

    # show
    for angle, img in zip(angles, imgs):
        fig, ax = plt.subplots()
        im = ax.imshow(img)
        fig.colorbar(im)
        fig.savefig(f"/Users/jhughes/Desktop/projection_test/{int(np.rad2deg(angle)):03d}.png")
        plt.close()
    #
    # model = np.zeros(densities.shape, dtype=np.float32)
    # for i, img in enumerate(imgs.reshape(len(xs), xs[0].shape[0], xs[0].shape[1])):
    #     model = backproject_3d(xs[i], ys[i], zs[i], img,
    #                                 model,
    #                                 mask, b, delta, norms[i], ds[i], True)
    # total = np.sum(img)
    # print(total, img.shape, imgs.shape)
    # model = ((model - total) / (len(imgs) - 1)).astype(np.float32)

    # # # 3d plot
    # fig = plt.figure()
    # ax = fig.add_subplot(projection='3d')
    #
    # n = 100
    #
    # for x, y, z in zip(xs, ys, zs):
    #     ax.scatter(x, y, z)
    #
    # # for norm, sx, sy in zip(norms, shift_xs, shift_ys):
    # #     print("TEST", norm)
    # #     ax.plot([0+sx, 100*norm[0]+sx], [0, 100*norm[1]], [0+sy, sy+100*norm[2]])
    #
    # ax.set_xlabel('X Label')
    # ax.set_ylabel('Y Label')
    # ax.set_zlabel('Z Label')
    #
    # plt.show()

    class Tomo(LinearOperator):
        def __init__(self, xs, ys, zs, norms, ds, b, delta, model_shape, mask, dtype=None):
            self.xs = xs
            self.ys = ys
            self.zs = zs
            self.norms = norms
            self.ds = ds
            self.b = b
            self.delta = delta
            self.model_shape = model_shape
            self.mask = mask
            super().__init__(dtype=np.dtype(dtype),
                             dims=self.model_shape,
                             dimsd=(len(self.xs), self.xs[0].shape[0], self.xs[0].shape[1]))

        def _matvec(self, densities):
            return np.array([project_3d(x, y, z, densities.reshape(self.model_shape).astype(np.float32),
                                        self.mask, self.b, self.delta, norm, d)
                    for x, y, z, norm, d in zip(self.xs, self.ys, self.zs, self.norms, self.ds)]).flatten()

        def _rmatvec(self, imgs):
            densitiesi = np.zeros(self.model_shape, dtype=np.float32)
            for i, img in enumerate(imgs.reshape(len(self.xs), self.xs[0].shape[0], self.xs[0].shape[1])):
                # densitiesi += backproject_3d(self.xs[i], self.ys[i], self.zs[i], img,
                #                densitiesi,
                #                self.mask, self.b, self.delta, self.norms[i], self.ds[i], True)
                densitiesi += backproject_3d(self.xs[i], self.ys[i], self.zs[i], img,
                                             np.zeros_like(densitiesi).astype(np.float32),
                                             self.mask, self.b, self.delta, self.norms[i], self.ds[i], True)
            # return ((densitiesi - np.sum(img)) / (len(self.xs)-1)).astype(np.float32)
            return densitiesi.flatten().astype(np.float32) #/ densitiesi.shape[0] / len(self.xs)
            #return densitiesi.flatten().astype(np.float32)

    print(len(xs))
    op = Tomo(xs, ys, zs, norms, ds, b, delta, densities.shape, mask, dtype=np.float32)

    from pylops.utils import dottest
    # # print(op.dims)
    _ = dottest(op, 140000, 1000000, atol=0.1, complexflag=0, verb=True)

    Dop = [
        pylops.FirstDerivative(
            (cube_size, cube_size, cube_size), axis=0, edge=False, kind="backward", dtype=np.float32
        ),
        pylops.FirstDerivative(
            (cube_size, cube_size, cube_size), axis=1, edge=False, kind="backward", dtype=np.float32
        ),
        pylops.FirstDerivative(
            (cube_size, cube_size, cube_size), axis=2, edge=False, kind="backward", dtype=np.float32
        )
    ]
    #
    # # TV
    # mu = 1.5
    # lamda = [0.1, 0.1, 0.1]
    # niter_out = 2
    # niter_in = 1
    #
    # model = pylops.optimization.sparsity.splitbregman(
    #     op,
    #     imgs.ravel(),
    #     Dop,
    #     niter_outer=niter_out,
    #     niter_inner=niter_in,
    #     mu=mu,
    #     epsRL1s=lamda,
    #     tol=1e-4,
    #     tau=1.0,
    #     show=True,
    #     **dict(iter_lim=5, damp=1e-4)
    # )[0]

    # model = op / imgs.flatten()
    # model = pylops.optimization.leastsquares.regularized_inversion(
    #     op, imgs.flatten(), Dop, **dict(iter_lim=10, show=True, atol=1E-8, btol=1E-8)
    # )[0]
    #
    # model = pylops.optimization.basic.lsqr(op, imgs.flatten(),  x0=np.random.rand(*densities.flatten().shape).astype(np.float32),
    #                                        niter=100, show=True, damp=0)[0]
    model = pylops.optimization.basic.lsqr(op, imgs.flatten(), niter=10, show=True)[0]#, x0 = np.random.randint(0, 30, densities.shape))[0] # x0=densities.flatten() + 50 * np.random.rand(*densities.flatten().shape).astype(np.float32) - 25)[0]
                                           #                                        niter=100, show=True, damp=0)[0]
    #model = pylops.optimization.basic.lsqr(op, imgs.flatten(), x0=densities.flatten(), niter=10, show=True)[0]
    # model = pylops.optimization.basic.lsqr(op, imgs.flatten(), niter=1, show=True)[0]

    # from pylops.optimization.cls_sparsity import FISTA
    #
    # fistasolver = FISTA(op)
    #
    # model = fistasolver.solve(imgs.flatten(), niter=2, show=True)[0]

    #model = pylops.optimization.basic.cgls(op, imgs.flatten(), niter=10, show=True)[0]
    model = model.reshape(densities.shape)

    limit = 150 # np.nanpercentile(model, 95)

    for i in range(cube_size):
        fig, axs = plt.subplots(ncols=2)
        im = axs[0].imshow(densities[:, :, i], vmin=0, vmax=150)#, vmin=0, vmax=5)
        fig.colorbar(im)
        im = axs[1].imshow(model[:, :, i], vmin=0, vmax=limit)
        fig.colorbar(im)
        fig.show()
        fig.savefig(f"/Users/jhughes/Desktop/projection_test/test_{i:03d}.png")

    # fig, axs = plt.subplots(ncols=2)
    # axs[0].imshow(densities[:, :, 2], vmin=0, vmax=10)
    # axs[1].imshow(model[:, :, 2], vmin=0, vmax=10)
    # fig.show()
    # fig.savefig("/Users/jhughes/Desktop/projection_test/reconstruction_0.png")
    #
    # fig, axs = plt.subplots(ncols=2)
    # axs[0].imshow(densities[:, :, 10], vmin=0, vmax=10)
    # axs[1].imshow(model[:, :, 10], vmin=0, vmax=10)
    # fig.show()
    # fig.savefig("/Users/jhughes/Desktop/projection_test/reconstruction_1.png")
    #
    #
    # fig, axs = plt.subplots(ncols=2)
    # axs[0].imshow(densities[:, :, 30], vmin=0, vmax=10)
    # axs[1].imshow(model[:, :, 30], vmin=0, vmax=10)
    # fig.show()
    # fig.savefig("/Users/jhughes/Desktop/projection_test/reconstruction_2.png")
    #
    #
    # fig, axs = plt.subplots(ncols=2)
    # axs[0].imshow(densities[:, :, 40], vmin=0, vmax=10)
    # axs[1].imshow(model[:, :, 40], vmin=0, vmax=10)
    # fig.show()
    # fig.savefig("/Users/jhughes/Desktop/projection_test/reconstruction_3.png")

    end_time = datetime.now()
    print(end_time - start_time)

import matplotlib.pyplot as plt
import numpy as np
import pylops

from tomograpy import project_3d, backproject_3d

def z_rotation_matrix_3d(angle):
    return np.array([[[np.cos(angle), np.sin(angle), 0],
                      [-np.sin(angle), np.cos(angle), 0],
                      [0, 0, 1]]])

angles = np.linspace(0, np.pi, 14) + np.pi/60

radius = 300

img_size = 100
cube_size = 100
# x, y = np.meshgrid(np.arange(img_size, dtype=np.float32), np.arange(img_size, dtype=np.float32))
# z = np.zeros((img_size, img_size), dtype=np.float32)
densities = np.zeros((cube_size, cube_size, cube_size), dtype=np.float32)
# densities[1:30, 1:30, 1:30] = 25
# densities[80:90, 80:90, 80:90] = 5
densities[30:-30, 30:-30, 30:-30] = 100
densities[40:-40, 40:-40, 40:-40] = 0

# densities[2, 2, 2] = 5
# densities[8, 8, 2] = 8
# densities[10:30, 10:30, 0:30] = 3
mask = np.ones((cube_size, cube_size, cube_size), dtype=bool)

b = (0, 0, 0)
b = (-cube_size / 2, -cube_size / 2, -cube_size / 2)

delta = (1.0, 1.0, 1.0)
# unit_normal = (1E-7, 1E-7, 1.0)
path_distance = 500.0

total = np.zeros_like(densities)

for angle in angles:
    t_angle = -angle + np.pi / 2  # np.deg2rad(np.abs(90*np.cos(angle)))

    img_x = np.arange(img_size) - img_size / 2
    img_y = np.zeros((img_size, img_size))
    img_z = np.arange(img_size) - img_size / 2
    img_x, img_z = np.meshgrid(img_x, img_z)

    img_x, img_y, img_z = img_x.flatten(), img_y.flatten(), img_z.flatten()

    R = z_rotation_matrix_3d(t_angle)
    coords = (R @ np.stack([img_x, img_y, img_z]))[0]
    img_x, img_y, img_z = coords[0], coords[1], coords[2]
    img_x = radius * np.cos(angle) + img_x  # - img_size/2
    img_y = radius * np.sin(angle) + img_y  # - img_size/2

    xx = img_x.reshape((img_size, img_size)).astype(np.float32)
    yy = -img_y.reshape((img_size, img_size)).astype(np.float32)
    zz = img_z.reshape((img_size, img_size)).astype(np.float32)

    v1 = np.array([xx[0, 1] - xx[0, 0], yy[0, 1] - yy[0, 0], zz[0, 1] - zz[0, 0]])
    v2 = np.array([xx[1, 0] - xx[0, 0], yy[1, 0] - yy[0, 0], zz[1, 0] - zz[0, 0]])
    v1 = v1 / np.linalg.norm(v1)
    v2 = v2 / np.linalg.norm(v2)
    normal = np.cross(v1, v2)
    normal = normal / np.linalg.norm(normal)

    norm = normal  # R @ np.array([1E-7, 1E-7, -1.0])
    norm[norm == 0] = 1E-6
    print(np.rad2deg(angle), np.rad2deg(t_angle), norm)

    proj = project_3d(xx, yy, zz, densities, mask, b, delta, norm, path_distance)

    fig, ax = plt.subplots()
    im = ax.imshow(proj, origin='lower')
    fig.colorbar(im)
    fig.savefig(f"/Users/jhughes/Desktop/projection_bpj/proj_{angle:0.3f}.png")
    plt.close()

    result = backproject_3d(xx, yy, zz, proj.astype(np.float32),
                       np.zeros_like(densities, dtype=np.float32),
                       mask, b, delta, norm, path_distance, True)
    print(result.shape)

    scaling = densities.shape[0]
    result = result / scaling

    total += result
# ((result - np.sum(proj)) / (len(self.xs)-1)).astype(np.float32)
#result = result - np.sum(proj)

# # Scale result to make sure that fbp(A, A(x)) == x holds at least
# # to some approximation. In limited experiments, this is true for
# # this version of FBP up to 1%.
# # *Note*: For some reason, we do not have to scale with respect to
# # the pixel dimension that is orthogonal to the rotation axis (`u`
# # or horizontal pixel dimension). Hence, we only scale with the
# # other pixel dimension (`v` or vertical pixel dimension).
# vg, pg = A.astra_compat_vg, A.astra_compat_pg
#
# pixel_height = (pg.det_size[0] / pg.det_shape[0])
# voxel_volume = np.prod(np.array(vg.size / np.array(vg.shape)))
# scaling = (np.pi / pg.num_angles) * pixel_height / voxel_volume
#
# rec *= scaling
# pixel_height = proj.size / proj.shape[0]
# voxel_volume = np.prod(densities.size / np.array(densities.shape))
# scaling = (np.pi / 1) * pixel_height / voxel_volume

total /= len(angles)

for i in range(cube_size):
    fig, axs = plt.subplots(ncols=2)
    im = axs[0].imshow(densities[i, :, :], origin='lower', vmin=0, vmax=150)
    fig.colorbar(im)

    im = axs[1].imshow(total[i, :, :], origin='lower', vmin=0, vmax=150)
    fig.colorbar(im)
    fig.savefig(f"/Users/jhughes/Desktop/projection_bpj/{i:03d}.png")
    plt.close()

import matplotlib.pyplot as plt
import numpy as np
import pylops

from tomograpy import project_3d, backproject_3d

def z_rotation_matrix_3d(angle):
    return np.array([[[np.cos(angle), np.sin(angle), 0],
                      [-np.sin(angle), np.cos(angle), 0],
                      [0, 0, 1]]])

angle =  np.pi/2
radius = 100

img_size = 100
cube_size = 100
# x, y = np.meshgrid(np.arange(img_size, dtype=np.float32), np.arange(img_size, dtype=np.float32))
# z = np.zeros((img_size, img_size), dtype=np.float32)
densities = np.zeros((cube_size, cube_size, cube_size), dtype=np.float32)
densities[1:30, 1:30, 1:30] = 25
densities[80:90, 80:90, 80:90] = 5

# densities[2, 2, 2] = 5
# densities[8, 8, 2] = 8
# densities[10:30, 10:30, 0:30] = 3
mask = np.ones((cube_size, cube_size, cube_size), dtype=bool)

b = (0, 0, 0)
b = (-cube_size / 2, -cube_size / 2, -cube_size / 2)

delta = (1.0, 1.0, 1.0)
#unit_normal = (1E-7, 1E-7, 1.0)
path_distance = 500.0

t_angle = -angle + np.pi / 2  # np.deg2rad(np.abs(90*np.cos(angle)))

img_x = np.arange(img_size) - img_size / 2
img_y = np.zeros((img_size, img_size))
img_z = np.arange(img_size) - img_size / 2
img_x, img_z = np.meshgrid(img_x, img_z)

img_x, img_y, img_z = img_x.flatten(), img_y.flatten(), img_z.flatten()

R = z_rotation_matrix_3d(t_angle)
coords = (R @ np.stack([img_x, img_y, img_z]))[0]
img_x, img_y, img_z = coords[0], coords[1], coords[2]
img_x = radius * np.cos(angle) + img_x  # - img_size/2
img_y = radius * np.sin(angle) + img_y  # - img_size/2

xx = img_x.reshape((img_size, img_size)).astype(np.float32)
yy = -img_y.reshape((img_size, img_size)).astype(np.float32)
zz = img_z.reshape((img_size, img_size)).astype(np.float32)

v1 = np.array([xx[0, 1] - xx[0, 0], yy[0, 1] - yy[0, 0], zz[0, 1] - zz[0, 0]])
v2 = np.array([xx[1, 0] - xx[0, 0], yy[1, 0] - yy[0, 0], zz[1, 0] - zz[0, 0]])
v1 = v1 / np.linalg.norm(v1)
v2 = v2 / np.linalg.norm(v2)
normal = np.cross(v1, v2)
normal = normal / np.linalg.norm(normal)

norm = normal  # R @ np.array([1E-7, 1E-7, -1.0])
norm[norm == 0] = 1E-6
print(np.rad2deg(angle), np.rad2deg(t_angle), norm)

proj = project_3d(xx, yy, zz, densities, mask, b, delta, norm, path_distance)

fig, ax = plt.subplots()
im = ax.imshow(proj, origin='lower')
fig.colorbar(im)
fig.savefig("/Users/jhughes/Desktop/projection_test2/proj.png")
plt.close()

result = backproject_3d(xx, yy, zz, proj.astype(np.float32),
                   np.zeros_like(densities, dtype=np.float32),
                   mask, b, delta, norm, path_distance, True)
print(result.shape)
# ((result - np.sum(proj)) / (len(self.xs)-1)).astype(np.float32)
#result = result - np.sum(proj)

# # Scale result to make sure that fbp(A, A(x)) == x holds at least
# # to some approximation. In limited experiments, this is true for
# # this version of FBP up to 1%.
# # *Note*: For some reason, we do not have to scale with respect to
# # the pixel dimension that is orthogonal to the rotation axis (`u`
# # or horizontal pixel dimension). Hence, we only scale with the
# # other pixel dimension (`v` or vertical pixel dimension).
# vg, pg = A.astra_compat_vg, A.astra_compat_pg
#
# pixel_height = (pg.det_size[0] / pg.det_shape[0])
# voxel_volume = np.prod(np.array(vg.size / np.array(vg.shape)))
# scaling = (np.pi / pg.num_angles) * pixel_height / voxel_volume
#
# rec *= scaling
# pixel_height = proj.size / proj.shape[0]
# voxel_volume = np.prod(densities.size / np.array(densities.shape))
# scaling = (np.pi / 1) * pixel_height / voxel_volume
scaling = densities.shape[0]
result = result / scaling

for i in range(cube_size):
    fig, axs = plt.subplots(ncols=2)
    im = axs[0].imshow(densities[i, :, :], origin='lower')
    fig.colorbar(im)

    im = axs[1].imshow(result[i, :, :], origin='lower')
    fig.colorbar(im)
    fig.savefig(f"/Users/jhughes/Desktop/projection_test2/{i:03d}.png")
    plt.close()

import matplotlib.pyplot as plt
import numpy as np
import pylops
from datetime import datetime
from glob import glob
from astropy.io import fits
import os
from itertools import chain
from tomograpy import project_3d, backproject_3d
from pylops.basicoperators import FunctionOperator
from pylops import LinearOperator, lsqr
from skimage.transform import resize

def z_rotation_matrix_3d(angle):
    return np.array([[[np.cos(angle), np.sin(angle), 0],
                      [-np.sin(angle), np.cos(angle), 0],
                      [0, 0, 1]]])

if __name__ == "__main__":
    print("Test started")
    start_time = datetime.now()

    angles = list(chain(range(70, 80), range(130, 140)))
    rad_angles = np.deg2rad(angles)

    path = "/Users/jhughes/Desktop/data/synthetic COMPLETE data/1 degree/"
    filenames = glob(path + "*.fits")
    loaded_imgs = {np.deg2rad(angle): fits.open(os.path.join(path + f"comp_wl_284_ang_{angle}.fits"))[0].data for angle in angles}

    radius = 300

    img_size = 100
    cube_size = 100

    densities = np.zeros((cube_size, cube_size, cube_size), dtype=np.float32)
    mask = np.ones((cube_size, cube_size, cube_size), dtype=bool)

    b = (0, 0, 0)
    b = (-cube_size / 2, -cube_size / 2, -cube_size / 2)

    delta = (1.0, 1.0, 1.0)
    path_distance = 500.0

    norms, xs, ys, zs, ds, imgs = [], [], [], [], [], []
    #angles = np.array([0, 30])
    for angle in rad_angles:
        t_angle = -angle + np.pi/2 #np.deg2rad(np.abs(90*np.cos(angle)))

        img_x = np.arange(img_size) - img_size / 2
        img_y = np.zeros((img_size, img_size))
        img_z = np.arange(img_size) - img_size / 2
        img_x, img_z = np.meshgrid(img_x, img_z)

        img_x, img_y, img_z = img_x.flatten(), img_y.flatten(), img_z.flatten()

        R = z_rotation_matrix_3d(t_angle)
        coords = (R @ np.stack([img_x, img_y, img_z]))[0]
        img_x, img_y, img_z = coords[0], coords[1], coords[2]
        img_x = radius * np.cos(angle) + img_x  # - img_size/2
        img_y = radius * np.sin(angle) + img_y  # - img_size/2

        xx = img_x.reshape((img_size, img_size)).astype(np.float32)
        yy = -img_y.reshape((img_size, img_size)).astype(np.float32)
        zz = img_z.reshape((img_size, img_size)).astype(np.float32)

        v1 = np.array([xx[0, 1] - xx[0, 0], yy[0, 1] - yy[0, 0], zz[0, 1] - zz[0, 0]])
        v2 = np.array([xx[1, 0] - xx[0, 0], yy[1, 0] - yy[0, 0], zz[1, 0] - zz[0, 0]])
        v1 = v1 / np.linalg.norm(v1)
        v2 = v2 / np.linalg.norm(v2)
        normal = np.cross(v1, v2)
        normal = normal / np.linalg.norm(normal)

        norm = normal  # R @ np.array([1E-7, 1E-7, -1.0])
        norm[norm == 0] = 1E-6
        norms.append(norm)
        print(np.rad2deg(angle), np.rad2deg(t_angle), norm)

        d = 500
        img = resize(loaded_imgs[angle], (img_size, img_size))  #project_3d(xx, yy, zz, densities, mask, b, delta, norm, d)
        xs.append(xx)
        ys.append(yy)
        zs.append(zz)
        ds.append(d)
        imgs.append(img.astype(np.float32))
    imgs = np.array(imgs)

    # show
    for angle, img in zip(rad_angles, imgs):
        fig, ax = plt.subplots()
        im = ax.imshow(img)
        fig.colorbar(im)
        fig.savefig(f"/Users/jhughes/Desktop/projection_solar/{int(np.rad2deg(angle)):03d}.png")
        plt.close()
    #
    # model = np.zeros(densities.shape, dtype=np.float32)
    # for i, img in enumerate(imgs.reshape(len(xs), xs[0].shape[0], xs[0].shape[1])):
    #     model = backproject_3d(xs[i], ys[i], zs[i], img,
    #                                 model,
    #                                 mask, b, delta, norms[i], ds[i], True)
    # total = np.sum(img)
    # print(total, img.shape, imgs.shape)
    # model = ((model - total) / (len(imgs) - 1)).astype(np.float32)

    # # # 3d plot
    # fig = plt.figure()
    # ax = fig.add_subplot(projection='3d')
    #
    # n = 100
    #
    # for x, y, z in zip(xs, ys, zs):
    #     ax.scatter(x, y, z)
    #
    # # for norm, sx, sy in zip(norms, shift_xs, shift_ys):
    # #     print("TEST", norm)
    # #     ax.plot([0+sx, 100*norm[0]+sx], [0, 100*norm[1]], [0+sy, sy+100*norm[2]])
    #
    # ax.set_xlabel('X Label')
    # ax.set_ylabel('Y Label')
    # ax.set_zlabel('Z Label')
    #
    # plt.show()

    class Tomo(LinearOperator):
        def __init__(self, xs, ys, zs, norms, ds, b, delta, model_shape, mask, dtype=None):
            self.xs = xs
            self.ys = ys
            self.zs = zs
            self.norms = norms
            self.ds = ds
            self.b = b
            self.delta = delta
            self.model_shape = model_shape
            self.mask = mask
            super().__init__(dtype=np.dtype(dtype),
                             dims=self.model_shape,
                             dimsd=(len(self.xs), self.xs[0].shape[0], self.xs[0].shape[1]))

        def _matvec(self, densities):
            return np.array([project_3d(x, y, z, densities.reshape(self.model_shape).astype(np.float32),
                                        self.mask, self.b, self.delta, norm, d)
                    for x, y, z, norm, d in zip(self.xs, self.ys, self.zs, self.norms, self.ds)]).flatten()

        def _rmatvec(self, imgs):
            densitiesi = np.zeros(self.model_shape, dtype=np.float32)
            for i, img in enumerate(imgs.reshape(len(self.xs), self.xs[0].shape[0], self.xs[0].shape[1])):
                # densitiesi += backproject_3d(self.xs[i], self.ys[i], self.zs[i], img,
                #                densitiesi,
                #                self.mask, self.b, self.delta, self.norms[i], self.ds[i], True)
                densitiesi += backproject_3d(self.xs[i], self.ys[i], self.zs[i], img,
                                             np.zeros_like(densitiesi).astype(np.float32),
                                             self.mask, self.b, self.delta, self.norms[i], self.ds[i], True)
            # return ((densitiesi - np.sum(img)) / (len(self.xs)-1)).astype(np.float32)
            return densitiesi.flatten().astype(np.float32) #/ densitiesi.shape[0] / len(self.xs)
            #return densitiesi.flatten().astype(np.float32)

    print(len(xs))
    op = Tomo(xs, ys, zs, norms, ds, b, delta, densities.shape, mask, dtype=np.float32)

    Dop = [
        pylops.FirstDerivative(
            (cube_size, cube_size, cube_size), axis=0, edge=False, kind="backward", dtype=np.float32
        ),
        pylops.FirstDerivative(
            (cube_size, cube_size, cube_size), axis=1, edge=False, kind="backward", dtype=np.float32
        ),
        pylops.FirstDerivative(
            (cube_size, cube_size, cube_size), axis=2, edge=False, kind="backward", dtype=np.float32
        )
    ]
    #
    # # TV
    # mu = 1.5
    # lamda = [0.1, 0.1, 0.1]
    # niter_out = 2
    # niter_in = 1
    #
    # model = pylops.optimization.sparsity.splitbregman(
    #     op,
    #     imgs.ravel(),
    #     Dop,
    #     niter_outer=niter_out,
    #     niter_inner=niter_in,
    #     mu=mu,
    #     epsRL1s=lamda,
    #     tol=1e-4,
    #     tau=1.0,
    #     show=True,
    #     **dict(iter_lim=5, damp=1e-4)
    # )[0]

    # model = op / imgs.flatten()
    # model = pylops.optimization.leastsquares.regularized_inversion(
    #     op, imgs.flatten(), Dop, **dict(iter_lim=10, show=True, atol=1E-8, btol=1E-8)
    # )[0]
    #
    # model = pylops.optimization.basic.lsqr(op, imgs.flatten(),  x0=np.random.rand(*densities.flatten().shape).astype(np.float32),
    #                                        niter=100, show=True, damp=0)[0]
    model = pylops.optimization.basic.lsqr(op, imgs.flatten(), niter=3, show=True)[0]#, x0 = np.random.randint(0, 30, densities.shape))[0] # x0=densities.flatten() + 50 * np.random.rand(*densities.flatten().shape).astype(np.float32) - 25)[0]
                                           #                                        niter=100, show=True, damp=0)[0]
    #model = pylops.optimization.basic.lsqr(op, imgs.flatten(), x0=densities.flatten(), niter=10, show=True)[0]
    # model = pylops.optimization.basic.lsqr(op, imgs.flatten(), niter=1, show=True)[0]

    # from pylops.optimization.cls_sparsity import FISTA
    #
    # fistasolver = FISTA(op)
    #
    # model = fistasolver.solve(imgs.flatten(), niter=2, show=True)[0]

    #model = pylops.optimization.basic.cgls(op, imgs.flatten(), niter=10, show=True)[0]
    model = model.reshape(densities.shape)

    np.save("/Users/jhughes/Desktop/projection_solar/cube.npy", model)
    limit = np.nanpercentile(model, 95)

    for i in range(cube_size):
        fig, ax = plt.subplots()
        im = ax.imshow(model[:, :, i], vmin=0, vmax=limit)
        fig.colorbar(im)
        fig.show()
        fig.savefig(f"/Users/jhughes/Desktop/projection_solar/test_{i:03d}.png")

    reconstructions = op @ model.flatten()
    reconstructions = reconstructions.reshape((len(xs), img_size, img_size))

    for i, angle in enumerate(rad_angles):
        fig, axs = plt.subplots(ncols=2)
        im = axs[0].imshow(imgs[i])
        fig.colorbar(im)
        im = axs[1].imshow(reconstructions[i])
        fig.colorbar(im)
        fig.show()
        fig.savefig(f"/Users/jhughes/Desktop/projection_solar/recon_{np.rad2deg(angle)}.png")

    end_time = datetime.now()
    print(end_time - start_time)

from .tomograpy import *
from . import third

__doc__ = tomograpy.__doc__
if hasattr(tomograpy, "__all__"):
    __all__ = tomograpy.__all__

import numpy as np

def print_hello(name):
    print(f"hello {name}")