import unittest

import numpy as np
import numpy.testing as npt
from scipy.ndimage import gaussian_filter

import sigpy as sp
from sigpy.mri import linop, rf, sim

if __name__ == "__main__":
    unittest.main()


class TestPtx(unittest.TestCase):
    @staticmethod
    def problem_2d(dim):
        img_shape = [dim, dim]
        sens_shape = [8, dim, dim]

        # target - slightly blurred circle
        x, y = np.ogrid[
            -img_shape[0] / 2 : img_shape[0] - img_shape[0] / 2,
            -img_shape[1] / 2 : img_shape[1] - img_shape[1] / 2,
        ]
        circle = x * x + y * y <= int(img_shape[0] / 6) ** 2
        target = np.zeros(img_shape)
        target[circle] = 1
        target = gaussian_filter(target, 1)
        target = target.astype(np.complex64)

        sens = sim.birdcage_maps(sens_shape)

        return target, sens

    @staticmethod
    def problem_3d(dim, Nz):
        Nc = 8
        img_shape = [dim, dim, Nz]
        sens_shape = [Nc, dim, dim, Nz]

        # target - slightly blurred circle
        x, y, z = np.ogrid[
            -img_shape[0] / 2 : img_shape[0] - img_shape[0] / 2,
            -img_shape[1] / 2 : img_shape[1] - img_shape[1] / 2,
            -img_shape[2] / 2 : img_shape[2] - img_shape[2] / 2,
        ]
        circle = x * x + y * y + z * z <= int(img_shape[0] / 5) ** 2
        target = np.zeros(img_shape)
        target[circle] = 1
        target = gaussian_filter(target, 1)
        target = target.astype(np.complex64)
        sens = sp.mri.sim.birdcage_maps(sens_shape)

        return target, sens

    def test_stspa_radial(self):
        target, sens = self.problem_2d(8)

        # makes dim*dim*2 trajectory
        traj = sp.mri.radial(
            (sens.shape[1], sens.shape[1], 2),
            target.shape,
            golden=True,
            dtype=np.float32,
        )
        # reshape to be Nt*2 trajectory
        traj = np.reshape(traj, [traj.shape[0] * traj.shape[1], 2])

        A = linop.Sense(sens, coord=traj, weights=None, ishape=target.shape).H

        pulses = rf.stspa(
            target,
            sens,
            traj,
            dt=4e-6,
            alpha=1,
            b0=None,
            st=None,
            explicit=False,
            max_iter=100,
            tol=1e-4,
        )

        npt.assert_array_almost_equal(A * pulses, target, 1e-3)

    def test_stspa_spiral(self):
        target, sens = self.problem_2d(8)

        fov = 0.55
        gts = 6.4e-6
        gslew = 190
        gamp = 40
        R = 1
        dx = 0.025  # in m
        # construct a trajectory
        g, k, t, s = rf.spiral_arch(fov / R, dx, gts, gslew, gamp)

        A = linop.Sense(sens, coord=k, ishape=target.shape).H

        pulses = rf.stspa(
            target,
            sens,
            k,
            dt=4e-6,
            alpha=1,
            b0=None,
            st=None,
            explicit=False,
            max_iter=100,
            tol=1e-4,
        )

        npt.assert_array_almost_equal(A * pulses, target, 1e-3)

    def test_stspa_2d_explicit(self):
        target, sens = self.problem_2d(8)
        dim = target.shape[0]
        g, k1, t, s = rf.spiral_arch(0.24, dim, 4e-6, 200, 0.035)
        k1 = k1 / dim

        A = rf.PtxSpatialExplicit(
            sens, k1, dt=4e-6, img_shape=target.shape, b0=None
        )
        pulses = sp.mri.rf.stspa(
            target,
            sens,
            st=None,
            coord=k1,
            dt=4e-6,
            max_iter=100,
            alpha=10,
            tol=1e-4,
            phase_update_interval=200,
            explicit=True,
        )

        npt.assert_array_almost_equal(A * pulses, target, 1e-3)

    def test_stspa_3d_explicit(self):
        nz = 4
        target, sens = self.problem_3d(3, nz)
        dim = target.shape[0]

        g, k1, t, s = rf.spiral_arch(0.24, dim, 4e-6, 200, 0.035)
        k1 = k1 / dim

        k1 = rf.stack_of(k1, nz, 0.1)
        A = rf.linop.PtxSpatialExplicit(
            sens, k1, dt=4e-6, img_shape=target.shape, b0=None
        )

        pulses = sp.mri.rf.stspa(
            target,
            sens,
            st=None,
            coord=k1,
            dt=4e-6,
            max_iter=30,
            alpha=10,
            tol=1e-3,
            phase_update_interval=200,
            explicit=True,
        )

        npt.assert_array_almost_equal(A * pulses, target, 1e-3)

    def test_stspa_3d_nonexplicit(self):
        nz = 3
        target, sens = self.problem_3d(3, nz)
        dim = target.shape[0]

        g, k1, t, s = rf.spiral_arch(0.24, dim, 4e-6, 200, 0.035)
        k1 = k1 / dim

        k1 = rf.stack_of(k1, nz, 0.1)
        A = sp.mri.linop.Sense(
            sens, k1, weights=None, tseg=None, ishape=target.shape
        ).H

        pulses = sp.mri.rf.stspa(
            target,
            sens,
            st=None,
            coord=k1,
            dt=4e-6,
            max_iter=30,
            alpha=10,
            tol=1e-3,
            phase_update_interval=200,
            explicit=False,
        )

        npt.assert_array_almost_equal(A * pulses, target, 1e-3)

    def test_spokes(self):
        # spokes problem definition:
        dim = 20  # size of the b1 matrix loaded
        n_spokes = 5
        fov = 20  # cm
        dx_max = 2  # cm
        gts = 4e-6
        sl_thick = 5  # slice thickness, mm
        tbw = 4
        dgdtmax = 18000  # g/cm/s
        gmax = 2  # g/cm

        _, sens = self.problem_2d(dim)
        roi = np.zeros((dim, dim))
        radius = dim // 2
        cx, cy = dim // 2, dim // 2
        y, x = np.ogrid[-radius:radius, -radius:radius]
        index = x**2 + y**2 <= radius**2
        roi[cy - radius : cy + radius, cx - radius : cx + radius][index] = 1
        sens = sens * roi

        [pulses, g] = rf.stspk(
            roi,
            sens,
            n_spokes,
            fov,
            dx_max,
            gts,
            sl_thick,
            tbw,
            dgdtmax,
            gmax,
            alpha=1,
        )

        # should give the number of pulses corresponding to number of TX ch
        npt.assert_equal(np.shape(pulses)[0], np.shape(sens)[0])
        # should hit the max gradient constraint
        npt.assert_almost_equal(gmax, np.max(g), decimal=3)
