import numpy as np
from numpy.polynomial.polynomial import Polynomial
# import timeit
import cProfile
import pstats

DISTANCE = True
RAYTRACE = True

def get_sun(phi, theta):
     return np.array([
            (r_maj - r_maj*np.cos(sun_r[0]))*np.cos(sun_r[1]),
            (r_maj - r_maj*np.cos(sun_r[0]))*np.sin(sun_r[1]),
            r_maj * np.sin(sun_r[0]),
        ])

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 uproot(arr):
    mask = np.logical_and(arr > 0, np.isreal(arr))
    masked = arr[mask]
    return np.real(sorted(masked))


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


def illumination(phi, theta, sun):

    def normal_vector(phi, theta):
        rx, ry, rz = surface_point(phi, theta)
        cx = r_maj * np.cos(theta)
        cy = 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(phi, theta):
        rx, ry, rz = surface_point(phi, theta)
        sx, sy, sz = sun
        retx, rety, retz = rx-sx, ry-sy, rz-sz
        length = np.linalg.norm([retx, rety, retz])
        norm = length*length if DISTANCE else length
        return np.array([retx, rety, retz])/norm

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

    rt = is_illuminated(phi, theta) if RAYTRACE else 1
    return np.dot(ray_vector(phi, theta), -normal_vector(phi, theta)) * rt


def run_round_illumination(res=100):
    phi = np.linspace(-np.pi, np.pi, res)
    theta = np.linspace(-np.pi, np.pi, res)
    day_cycle = np.linspace(-np.pi, np.pi, 24)
    sun = get_sun(day_cycle[0], 0)
    illu = np.array([[illumination(ph, th, sun) for th in theta] for ph in phi])
    for sun_pos in day_cycle[1:]:
        sun = get_sun(sun_pos, 0)
        illu += np.array([[illumination(ph, th, sun) for th in theta] for ph in phi])
    illu = illu / np.amax(illu)

##########################################


def illumination2(phi, theta, sun):
    def sp(phi, theta):
        rx = (r_maj - r_min*np.cos(phi)) * np.cos(theta)
        ry = (r_maj - r_min*np.cos(phi)) * np.sin(theta)
        rz = r_min * np.sin(phi)
        return np.array([rx, ry, rz])

    def cp(phi, theta):
        cx = r_maj * np.cos(theta)
        cy = r_maj * np.sin(theta)
        cz = 0
        return np.array([cx, cy, cz])

    def rv(surf):
        ray = surf - sun
        length = np.linalg.norm(ray, axis=-1)
        normalizer = length*length if DISTANCE else length
        return ray/normalizer

    surf = surface_point(phi, theta)
    norm = surf - cp(phi, theta)
    ray = rv(surf)

    if RAYTRACE:
        pol = Polynomial(coef(r_maj, r_min, ray, sun))
        roots = uproot(pol.roots())
        if len(roots) == 0:
            return False
        rt = np.allclose(surface_point(phi, theta), sun+roots[0]*ray)
    else:
        rt = 1

    return np.dot(ray, -norm) * rt


def run_round_illumination2(res=100):
    phi = np.linspace(-np.pi, np.pi, res)
    theta = np.linspace(-np.pi, np.pi, res)
    day_cycle = np.linspace(-np.pi, np.pi, 24)
    sun = get_sun(day_cycle[0], 0)
    illu = np.array([[illumination2(ph, th, sun) for th in theta] for ph in phi])
    for sun_pos in day_cycle[1:]:
        sun = get_sun(sun_pos, 0)
        illu += np.array([[illumination2(ph, th, sun) for th in theta] for ph in phi])
    illu = illu / np.amax(illu)

#########################################


def illumination3(phi, theta, sun):
    def sp(phi, theta):
        rx = (r_maj - r_min*np.cos(phi)) * np.cos(theta)
        ry = (r_maj - r_min*np.cos(phi)) * np.sin(theta)
        rz = r_min * np.sin(phi)
        return np.array([rx, ry, rz])

    def cp(phi, theta):
        cx = r_maj * np.cos(theta)
        cy = r_maj * np.sin(theta)
        cz = 0
        return np.array([cx, cy, cz])

    def rv(surf):
        ray = surf - sun
        length = np.linalg.norm(ray, axis=-1)
        normalizer = length*length if DISTANCE else length
        return ray/normalizer

    surf = surface_point(phi, theta)
    norm = surf - cp(phi, theta)
    ray = rv(surf)

    if RAYTRACE:
        pol = Polynomial(coef(r_maj, r_min, ray, sun))
        roots = uproot(pol.roots())
        if len(roots) == 0:
            return False
        rt = np.allclose(surface_point(phi, theta), sun+roots[0]*ray)
    else:
        rt = 1

    return np.dot(ray, -norm) * rt


def run_illumination(n=100):
    phi = np.linspace(-np.pi, np.pi, n)
    theta = np.linspace(-np.pi, np.pi, n)
    np.array([[illumination(ph, th) for th in theta] for ph in phi])


def run_round_illumination3(res=100):
    # TODO: refactor TorusWorld.illumination() do use vectoization!!
    phi = np.linspace(-np.pi, np.pi, res)
    theta = np.linspace(-np.pi, np.pi, res)
    day_cycle = np.linspace(-np.pi, np.pi, 24)
    sun = get_sun(day_cycle[0], 0)
    illu = np.array(illumination2(phi, theta, sun))
    for sun_pos in day_cycle[1:]:
        sun = get_sun(sun_pos, 0)
        illu += np.array(illumination2(phi, theta, sun))
    illu = illu / np.amax(illu)


r_maj = 1
r_min = 0.5

sun_r = (np.random.rand(2)*2-1)*np.pi
sun = get_sun(*sun_r)

profiler = cProfile.Profile()
profiler.enable()
run_round_illumination()
profiler.disable()
stats = pstats.Stats(profiler.print_stats('cumtime'))
stats.strip_dirs().print_stats()