import numpy as np
from numpy.polynomial.polynomial import Polynomial
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from matplotlib.widgets import Slider, Button


def uproot(arr):
    mask = np.logical_and(arr > 0, np.isreal(arr))
    masked = arr[mask]
    return np.real(sorted(masked))


class TorusWorld:

    @staticmethod
    def coef(rmaj, rmin, d, o):
        """
            curtesy of: http://blog.marcinchwedczuk.pl/ray-tracing-torus
        """
        k1 = np.inner(d, d)
        k2 = np.inner(o, d)
        k3 = np.inner(o, o) - (rmin**2 + rmaj**2)

        c4 = k1**2
        c3 = 4*k1*k2
        c2 = 2*k1*k3 + 4*k2**2 + 4*(rmaj*d[2])**2
        c1 = 4*k3*k2 + 8*(rmaj**2)*o[2]*d[2]
        c0 = k3**2 - 4*(rmaj**2)*(rmin**2 - o[2]**2)
        return [c0, c1, c2, c3, c4]

    def __init__(self, rfrac):
        self.r_maj = 1
        self.r_min = rfrac
        self.sun = np.array([self.r_maj, 0, self.r_maj])
        self.sun_r = np.array([np.pi/2, 0])
        self.surface_map = None

    def update(self, rfrac=None):
        rfrac = rfrac if rfrac else self.r_min
        self.r_min = rfrac
        self.put_sun(*self.sun_r)

    def put_sun(self, phi, theta):
        self.sun_r = np.array([phi, theta])
        self.sun = np.array([
            (self.r_maj - self.r_maj*np.cos(phi))*np.cos(theta),
            (self.r_maj - self.r_maj*np.cos(phi))*np.sin(theta),
            self.r_maj * np.sin(phi),
        ])

    def surface_point(self, phi, theta):
        rx = (self.r_maj - self.r_min*np.cos(phi)) * np.cos(theta)
        ry = (self.r_maj - self.r_min*np.cos(phi)) * np.sin(theta)
        rz = self.r_min * np.sin(phi)
        return rx, ry, rz

    def normal_vector(self, phi, theta):
        rx, ry, rz = self.surface_point(phi, theta)
        cx = self.r_maj * np.cos(theta)
        cy = self.r_maj * np.sin(theta)
        cz = 0
        retx, rety, retz = rx-cx, ry-cy, rz-cz
        return np.array([retx, rety, retz])/np.linalg.norm([retx, rety, retz])

    def ray_vector(self, phi, theta):
        rx, ry, rz = self.surface_point(phi, theta)
        sx, sy, sz = self.sun
        retx, rety, retz = rx-sx, ry-sy, rz-sz
        return np.array([retx, rety, retz])/np.linalg.norm([retx, rety, retz])

    def is_illuminated(self, phi, theta):
        ray = self.ray_vector(phi, theta)
        pol = Polynomial(self.coef(self.r_maj, self.r_min, ray, self.sun))
        roots = uproot(pol.roots())
        if len(roots) == 0:
            return False
        return np.allclose(self.surface_point(phi, theta), self.sun+roots[0]*ray)

    def illumination(self, phi, theta):
        # return np.dot(self.ray_vector(phi, theta), -self.normal_vector(phi, theta)) 
        return np.dot(self.ray_vector(phi, theta), -self.normal_vector(phi, theta)) * self.is_illuminated(phi, theta)


class NoInamge():
    def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
        fig, (ax_side, ax_top, ax_map) = plt.subplots(3, 1, gridspec_kw={'height_ratios': [2, 2, 1]})
        self.fig = fig
        self.ax = dict(
            map=ax_map,
            top=ax_top,
            side=ax_side,
        )
        self.lines = dict()

        self.sun_kwargs = dict(marker='o', color='r', markersize=10,)
        self.contour_kwargs = dict(cmap=plt.colormaps['hot'], vmin=0, vmax=1)
        self.levels = 25  # [-1, 0, 1]

        self.torus = TorusWorld(rfrac_init)
        self.init_top_view()
        self.init_side_view()
        self.init_map_view()

    def init_map_view(self):
        self.lines['map_border'], = self.ax['map'].plot(*self._mantle_map(), 'k')
        self.lines['pos_map'], = self.ax['map'].plot(*self._sunpos_map(), marker='o', color='r', markersize=10,)
        self.lines['dawnline_map'] = [
            self.ax['map'].contourf(*self._contour_map(), self.levels, **self.contour_kwargs),
            ]
        self.ax['map'].set_aspect('equal')
        self.ax['map'].set_ylabel('Longitude')
        self.ax['map'].set_xlabel('Lattitude')
        self.ax['map'].axis('off')

    def init_top_view(self):
        self.lines['circles_top'], = self.ax['top'].plot(*self._top_section(), 'k')
        self.lines['path_top'], = self.ax['top'].plot(*self._sunpath_top(), 'k:')
        self.lines['pos_top'], = self.ax['top'].plot(*self._sunpos_top(), **self.sun_kwargs)
        self.ax['top'].set_xlim(-2.05, 2.05)
        self.ax['top'].set_ylim(-2.05, None)
        self.ax['top'].set_aspect('equal')

        self.ax['top'].axis('off')

    def init_side_view(self):
        self.lines['circles_side'], = self.ax['side'].plot(*self._crossection(), 'k')
        self.lines['path_side'], = self.ax['side'].plot(*self._sunpath_side(), 'k:')
        self.lines['pos_side'], = self.ax['side'].plot(*self._sunpos_side(), **self.sun_kwargs)
        self.ax['side'].set_xlim(-2.05, 2.05)
        self.ax['side'].set_aspect('equal')
        # self.ax['side'].set_ylim(-self.torus.r_min, None)
        self.ax['side'].axis('off')

    def _mantle_map(self, n=1000):
        phi = np.linspace(-np.pi, np.pi, n)
        u = phi * self.torus.r_min
        width = np.pi*(self.torus.r_maj - self.torus.r_min*np.cos(phi))
        data = np.array([width, u])
        data2 = np.array([-width, u[::-1]])

        data = np.append(data, data2, axis=1)
        data = np.append(data, data[:, 0:1], axis=1)
        return data

    def _contour_map(self, n=100):
        theta = np.linspace(-np.pi, np.pi, n)
        phi = np.linspace(-np.pi, np.pi, int(n*self.torus.r_min))

        y = np.array([[w]*n for w in self.torus.r_min*phi])
        func = (np.pi)*(self.torus.r_maj - self.torus.r_min*np.cos(phi))
        x = np.array([np.linspace(-f, f, n) for f in func])
        # x, y = np.meshgrid(np.linspace(-func[0], func[0], n), rmin*theta)

        z = np.array([[self.torus.illumination(ph, th) for th in theta] for ph in phi])

        return x, y, z

    def _sunpos_map(self):
        phi, theta = self.torus.sun_r
        x = (self.torus.r_maj - self.torus.r_min*np.cos(phi)) * theta
        y = self.torus.r_min * phi
        return x, y

    def _top_section(self, n=1000):
        phi = np.linspace(0, np.pi, n)
        x1 = (self.torus.r_maj + self.torus.r_min) * np.cos(phi)
        y1 = -1 * (self.torus.r_maj + self.torus.r_min) * np.sin(phi)
        data1 = np.array([x1, y1])
        x2 = (self.torus.r_maj - self.torus.r_min) * np.cos(phi)
        y2 = -1 * (self.torus.r_maj - self.torus.r_min) * np.sin(phi)
        data2 = np.array([x2, y2])

        data = np.append(data1, [[np.nan,], [np.nan,]], axis=1)
        data = np.append(data, data2, axis=1)
        return data

    def _sunpath_top(self, n=1000):
        phi = np.linspace(0, np.pi, n)
        x = self.torus.r_maj + self.torus.r_maj*np.cos(phi)
        y = np.zeros_like(phi)
        return np.array([x, y])

    def _sunpos_top(self):
        phi, theta = self.torus.sun_r
        x = self.torus.r_maj + self.torus.r_maj*np.cos(phi)
        y = 0
        return x, y

    def _crossection(self, n=1000):
        phi = np.linspace(0, 2*np.pi, n)
        x = self.torus.r_maj + self.torus.r_min*np.cos(phi)
        y = self.torus.r_min*np.sin(phi)
        data = np.array([x, y])
        data2 = np.array([-x, y])
        data = np.append(data, [[np.nan,], [np.nan,]], axis=1)
        data = np.append(data, data2, axis=1)
        return data

    def _sunpath_side(self, n=1000):
        phi = np.linspace(0, 2*np.pi, n)
        x = self.torus.r_maj + self.torus.r_maj*np.cos(phi)
        y = self.torus.r_maj*np.sin(phi)
        return np.array([x, y])

    def _sunpos_side(self):
        phi, theta = self.torus.sun_r
        x = self.torus.r_maj - self.torus.r_maj*np.cos(phi)
        y = self.torus.r_maj*np.sin(phi)
        return x, y

    @staticmethod
    def redraw_plot(ax, line, func):
        x, y = func()
        line.set_xdata(x)
        line.set_ydata(y)
        ax.relim()
        ax.autoscale_view()

    @staticmethod
    def redraw_contourf(ax, container, func, levels=None, contour_kwargs=None):
        levels = levels if levels else [-1, 0, 1]
        contour_kwargs = contour_kwargs if contour_kwargs else {}
        for coll in container[0].collections:
            coll.remove()
        container[0] = ax.contourf(*func(), levels, **contour_kwargs)

    def update_torus(self, rfrac):
        self.torus.update(rfrac)
        self.redraw_plot(self.ax['map'], self.lines['map_border'], self._mantle_map)
        self.redraw_plot(self.ax['side'], self.lines['circles_side'], self._crossection)
        self.redraw_plot(self.ax['side'], self.lines['path_side'], self._sunpath_side)
        self.redraw_plot(self.ax['top'], self.lines['circles_top'], self._top_section)
        self.redraw_plot(self.ax['top'], self.lines['path_top'], self._sunpath_top)
        self.redraw_contourf(self.ax['map'], self.lines['dawnline_map'],
                             self._contour_map, self.levels, self.contour_kwargs)
        self.fig.canvas.draw_idle()

    def update_sun(self, phi, theta):
        self.torus.put_sun(phi, theta)
        self.redraw_plot(self.ax['map'], self.lines['pos_map'], self._sunpos_map)
        self.redraw_plot(self.ax['side'], self.lines['pos_side'], self._sunpos_side)
        self.redraw_plot(self.ax['top'], self.lines['pos_top'], self._sunpos_top)
        self.redraw_contourf(self.ax['map'], self.lines['dawnline_map'],
                             self._contour_map, self.levels, self.contour_kwargs)
        self.fig.canvas.draw_idle()


class ImageStatic(NoInamge):

    def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
        super().__init__(rfrac_init, sun_init)
        plt.show()


class ImageInteractive(NoInamge):

    def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
        super().__init__(rfrac_init, sun_init)
        self.init_interactivity(rfrac_init, sun_init)
        plt.show()

    def init_interactivity(self, rfrac_init, sun_init):
        self.fig.subplots_adjust(left=0.25, bottom=0.25)
        self.ax['slider_sun'] = self.fig.add_axes([0.25, 0.1, 0.65, 0.03])
        self.ax['slider_rf'] = self.fig.add_axes([0.1, 0.25, 0.0225, 0.63])
        self.ax['button_reset'] = self.fig.add_axes([0.8, 0.025, 0.1, 0.04])
        self.sliders = dict(
            sun_phi=Slider(
                ax=self.ax['slider_sun'],
                label='Angle of Sun',
                valmin=-np.pi,
                valmax=np.pi,
                valinit=sun_init,
            ),
            rfrac=Slider(
                ax=self.ax['slider_rf'],
                label="Fraction of Radii (r/R)",
                valmin=0,
                valmax=1,
                valinit=rfrac_init,
                orientation="vertical"
            )

        )
        self.sliders['sun_phi'].on_changed(self._slider_update_sun)
        self.sliders['rfrac'].on_changed(self._slider_update_torus)
        button = Button(self.ax['button_reset'], 'Reset', hovercolor='0.975')
        button.on_clicked(self._reset)

    def _slider_update_torus(self, val):
        self.update_torus(val)

    def _slider_update_sun(self, val):
        self.update_sun(val, 0)

    def _reset(self, event):
        self.sliders['sun_phi'].reset()
        self.sliders['rfrac'].reset()


if __name__ == '__main__':
    # ImageStatic()
    ImageInteractive()