import numpy as np
import matplotlib.patches

# https://matplotlib.org/mpl_toolkits/mplot3d/tutorial.html

class Arrow3D(matplotlib.patches.FancyArrowPatch):
    def __init__(self, xs, ys, zs, *args, **kwargs):
        matplotlib.patches.FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
        self._verts3d = xs, ys, zs

    def draw(self, renderer):
        from mpl_toolkits.mplot3d import proj3d
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        matplotlib.patches.FancyArrowPatch.draw(self, renderer)


def pathpatch_2d_to_3d_affine(pathpatch, mat_rot=np.array([[1,0,0],[0,1,0],[0,0,1]]), vec_trans=np.array([0,0,0])):
    """
    Transforms a 2D Patch to a 3D patch using the affine tranform
    of the given rotation matrix and translation vector.
    The pathpatch is assumed to be on the plane Z = 0.
    """
    import mpl_toolkits.mplot3d.art3d as art3d

    path = pathpatch.get_path() #Get the path and the associated transform
    trans = pathpatch.get_patch_transform()

    path = trans.transform_path(path) #Apply the transform

    pathpatch.__class__ = art3d.PathPatch3D #Change the class
    pathpatch._code3d = path.codes #Copy the codes
    pathpatch._facecolor3d = pathpatch.get_facecolor #Get the face color    

    verts = path.vertices #Get the vertices in 2D

    M = np.array([
        [mat_rot[0, 0], mat_rot[0, 1], mat_rot[0, 2], vec_trans[0]],
        [mat_rot[1, 0], mat_rot[1, 1], mat_rot[1, 2], vec_trans[1]],
        [mat_rot[2, 0], mat_rot[2, 1], mat_rot[2, 2], vec_trans[2]],
    ])

    pathpatch._segment3d = np.array([np.dot(M, (x, y, 0, 1)) for x, y in verts])


def plot_geometry(geo, angle=0):
    """
    Plots the given geometry.
    """
    import mpl_toolkits.mplot3d.art3d as art3d
    import matplotlib.pyplot as plt

    fig = plt.figure(figsize=(6,6))
    fig.suptitle('Cone Beam Compute Tomography geometry')
    ax = fig.add_subplot(111, projection='3d')
    ax.set_title('Current CBCT geometry, in scale')

    limXY = max(geo.DSO, geo.DSD-geo.DSO)
    limZ = geo.sVoxel[0]

    ax.set_box_aspect((limXY,limXY,limZ))

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_xlim3d(-limXY*1.2, limXY*1.2)
    ax.set_ylim3d(-limXY*1.2, limXY*1.2)
    ax.set_zlim3d(-limZ *1.2, limZ *1.2)

    # Trajectory of Source
    # https://matplotlib.org/devdocs/api/_as_gen/matplotlib.patches.Circle.htm
    circ = matplotlib.patches.Circle((0,0), geo.DSO, color='black', fill=False, ls='-.', lw=0.5)
    ax.add_patch(circ)
    art3d.pathpatch_2d_to_3d(circ, z=0, zdir="z")

    # Trajectory of Detector
    circ = matplotlib.patches.Circle((0,0), geo.DSD-geo.DSO, color='black', fill=False, ls='-.', lw=0.5)
    ax.add_patch(circ)
    art3d.pathpatch_2d_to_3d(circ, z=0, zdir="z")

    # Source
    # ax.scatter([0], [0], [0], color="g", s=100)
    sourcePos3D = [geo.DSO*np.cos(angle), geo.DSO*np.sin(angle), 0] # xyz
    ax.scatter([sourcePos3D[0]], [sourcePos3D[1]], [sourcePos3D[2]], color="steelblue", s=100)

    # Axes XYZ
    length_axis = geo.sVoxel[0]
    x_axis = Arrow3D([0, length_axis], [0, 0], [0, 0], mutation_scale=10, shrinkA=0, lw=1, arrowstyle="-|>", color="r")
    y_axis = Arrow3D([0, 0], [0, length_axis], [0, 0], mutation_scale=10, shrinkA=0, lw=1, arrowstyle="-|>", color="b")
    z_axis = Arrow3D([0, 0], [0, 0], [0, length_axis], mutation_scale=10, shrinkA=0, lw=1, arrowstyle="-|>", color="g")
    ax.add_artist(x_axis)
    ax.add_artist(y_axis)
    ax.add_artist(z_axis)

    # Detector
    print("sDetector = {}".format(geo.sDetector))
    alpha_detector = 0.7
    detectorPos3D = [
                        (geo.DSD-geo.DSO)*np.cos(angle+np.pi) + geo.offDetector[1]*np.sin(angle+np.pi),
                        (geo.DSD-geo.DSO)*np.sin(angle+np.pi) - geo.offDetector[1]*np.cos(angle+np.pi),
                         geo.offDetector[0]] # xyz
    detector = matplotlib.patches.Rectangle(
                    (-geo.sDetector[1]/2, -geo.sDetector[0]/2),
                    geo.sDetector[1], geo.sDetector[0], 
                    angle=0,#angle, 
                    color='maroon', 
                    fill=True, 
                    alpha=alpha_detector)
    print("detector={}".format(detector))
    ax.add_patch(detector)
    mat_rot = np.array([
        [-np.sin(angle), 0, np.cos(angle)],
        [ np.cos(angle), 0, np.sin(angle)],
        [ 0            , 1, 0            ]
    ])
    pathpatch_2d_to_3d_affine(detector, mat_rot, detectorPos3D)

    # Image FOV
    alpha_img = 0.1
    offOrigin = np.array([geo.offOrigin[2], geo.offOrigin[1], geo.offOrigin[0]])
    mat_rot_xy = np.array([
        [ 1, 0, 0],
        [ 0, 1, 0],
        [ 0, 0, 1]
    ], dtype=np.float)           
    for idx in range(2):
        img_face_xy = matplotlib.patches.Rectangle(
                        (-geo.sVoxel[2]/2, -geo.sVoxel[1]/2), # xy order
                        geo.sVoxel[2], geo.sVoxel[1], 
                        angle=0,#angle, 
                        color='black', 
                        fill=True, 
                        alpha=alpha_img)
        ax.add_patch(img_face_xy)
        pathpatch_2d_to_3d_affine(img_face_xy, mat_rot_xy, np.array([0, 0, -geo.sVoxel[0]/2 if idx == 0 else geo.sVoxel[0]/2])+offOrigin)

    mat_rot_yz = np.array([
        [ 0, 0, 1],
        [ 1, 0, 0],
        [ 0, 1, 0]
    ], dtype=np.float)           
    for idx in range(2):
        img_face_yz = matplotlib.patches.Rectangle(
                        (-geo.sVoxel[1]/2, -geo.sVoxel[0]/2), # xy order
                        geo.sVoxel[1], geo.sVoxel[0], 
                        angle=0,#angle, 
                        color='black', 
                        fill=True, 
                        alpha=alpha_img)
        ax.add_patch(img_face_yz)
        pathpatch_2d_to_3d_affine(img_face_yz, mat_rot_yz, np.array([-geo.sVoxel[2]/2 if idx == 0 else geo.sVoxel[2]/2, 0, 0])+offOrigin)

    mat_rot_zx = np.array([
        [ 0, 1, 0],
        [ 0, 0, 1],
        [ 1, 0, 0]
    ], dtype=np.float)           
    for idx in range(2):
        img_face_zx = matplotlib.patches.Rectangle(
                        (-geo.sVoxel[0]/2, -geo.sVoxel[2]/2), # xy order
                        geo.sVoxel[0], geo.sVoxel[2], 
                        angle=0,#angle, 
                        color='black', 
                        fill=True, 
                        alpha=alpha_img)
        ax.add_patch(img_face_zx)
        pathpatch_2d_to_3d_affine(img_face_zx, mat_rot_zx, np.array([0, -geo.sVoxel[1]/2 if idx == 0 else geo.sVoxel[1]/2, 0])+offOrigin)

    plt.show()


if __name__ == '__main__':
    nVoxelZ = 128
    nVoxelY = 256
    nVoxelX = 512
    nDetectorU = 400
    nDetectorV = 300
    offDetectorU = 0#500
    offDetectorV = 0#500
    offOriginX = 0#300
    offOriginY = 0#100
    offOriginZ = 0#100
    geo = tigre.geometry(mode='cone',default=True)
    geo.nVoxel = np.array([nVoxelZ,nVoxelY,nVoxelX])
    geo.sVoxel = geo.nVoxel
    geo.dVoxel = geo.sVoxel / geo.nVoxel
    geo.nDetector = np.array([nDetectorV,nDetectorU])
    geo.sDetector = geo.nDetector*geo.dDetector
    geo.offDetector = np.array([offDetectorV, offDetectorU])
    geo.offOrigin = np.array([offOriginZ, offOriginY, offOriginX])
    print(geo)
    angle = -np.pi/6
    plot_geometry(geo, angle)
