import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from pynufft import NUFFT
import pkg_resources
import scipy.misc
from OCTFrames import FrameManager,cachedOCT
from octReader import OCTManager
import scipy.ndimage
import scipy as S
from matplotlib.widgets import Slider
import attr

@attr.s(kw_only=True)
class RadialVolume:
    angle:int = attr.ib(default=2*np.pi)
    scan_count:int = attr.ib(default=100)
    scan_depth:int = attr.ib(default=800)
    diameter:int = attr.ib(default=1000)

    def set_data(self,data):
        radius = int(self.diameter/2)
        volume = data[0,:,:,slice(0,self.scan_depth)]
        volume = np.transpose(volume, (1, 0, 2))
        vol1 = np.flip(volume[:radius,:,:],axis=0)
        vol2 = volume[radius:,:,:]
        vol = np.hstack([vol1,vol2])
        self.vol = vol
    def get_coords(self) -> np.ndarray:
        om = np.meshgrid(np.linspace(0,self.angle,self.scan_count*2), np.arange(0,int(self.diameter/2),1)) #rectangular plot of polar data
        theta = np.ravel(om[0])
        r = np.ravel(om[1])
        om = np.empty((len(theta),2), dtype = np.float32)
        om[:,0] = r
        om[:,1] = theta
        return om
        # diameter = self.diameter
        # radius = int(diameter/2)
        # spoke_range = np.arange(0, radius)  # normalized between 0 and pi
        # M = radius*self.scan_count*2
        # om = np.empty((M,2), dtype = np.float32)
        # for i,angle in enumerate(np.linspace(0,360,self.scan_count*2)):
        #     radian = angle * 2 * np.pi/ 360.0
        #     spoke_x =  spoke_range * np.cos(radian)
        #     spoke_y =  spoke_range * np.sin(radian)
        #     om[radius*i : radius*(i + 1) ,0] = spoke_x
        #     om[radius*i : radius*(i + 1) ,1] = spoke_y
        # return om

    def get_bscan(self,phi=0):
        ind2 = int(self.scan_count+phi)
        d1 = np.flip(self.vol[:,0,:],axis=0)
        d2 = self.vol[:,ind2,:]
        frame = np.vstack([d1,d2])
        return frame
    def plot_bscan(self,phi=0):
        fig, ax = plt.subplots()
        frame = self.get_bscan(phi=phi)
        ax.imshow(frame, cmap=mpl.cm.gray,aspect='auto')
        fig.set_size_inches(12,5,forward=True)
        plt.show()
    def plot_vol(self):
        theta,rad = np.meshgrid(np.linspace(0,self.angle,self.scan_count*2), np.arange(0,int(self.diameter/2),1)) #rectangular plot of polar data
        X = theta
        Y = rad
        fig = plt.figure()
        ax3 = fig.add_subplot(131)
        ax = fig.add_subplot(132, polar=True)
        ax2 = fig.add_subplot(133)
        ax3.imshow(self.get_bscan(phi=0),cmap=mpl.cm.gray, aspect='auto')
        z_line, = ax3.plot([0,0],[0,self.diameter],color='r',linewidth=2)
        ax.pcolormesh(X, Y, self.vol[:,:,0],cmap=mpl.cm.gray, shading='auto') #X,Y & data2D must all be same dimensions
        ax2.imshow(self.vol[:,:,0],cmap=mpl.cm.gray, aspect='auto')
        axcolor = 'lightgoldenrodyellow'
        axpos = plt.axes([0.2, 0.1, 0.65, 0.03], facecolor=axcolor)
        spos = Slider(axpos, 'depth', 0, self.scan_depth,valinit=0,valstep=1)

        def update(val):
            z = spos.val
            frame = self.vol[:,:,z]
            z_line.set_xdata([z,z])
            ax2.imshow(frame,cmap=mpl.cm.gray,aspect='auto')
            ax.pcolormesh(X, Y, frame,cmap=mpl.cm.gray, shading='auto')
            fig.canvas.draw_idle()

        spos.on_changed(update)
        plt.show()

    def plot(self):
        pass

# Initiate the NufftObj object
example_radial =  'ExampleFiles/6977_OS_R_4_0_0000038.OCT'
# example_radial =  'ExampleFiles/03727_OD_R_0_0_0000030.OCT'
f = cachedOCT(example_radial)
oct = OCTManager(f)
oct.load()
volumes, angles, scan_diam, depth = oct.volume.shape

vol = RadialVolume(angle=2*np.pi,scan_count = angles, scan_depth = depth, diameter = scan_diam)
vol.set_data(oct.volume)
om = np.array(vol.get_coords())

Nd = (scan_diam,scan_diam)
Kd = (scan_diam*2,scan_diam*2)
Jd = (3,3)
NufftObj = NUFFT()
NufftObj.plan(om, Nd, Kd, Jd)
y = np.ravel(vol.vol[:,:,250])
image = NufftObj.solve(y, solver='L1TVOLS',maxiter=50,rho=0.1)
# image = NufftObj.adjoint(y)
fig,ax = plt.subplots()
ax.imshow(image.real,cmap=mpl.cm.gray,aspect='auto',norm=mpl.colors.Normalize(vmin=0.0, vmax=1))
plt.show()
vol.plot_vol()
