import numpy as np
import math
from PIL import Image
import time
import threading
import random
from multiprocessing import Pool

zz = 2 ** 9
aspect_ratio = 16.0 / 9.0
height = zz
width = math.floor(zz * aspect_ratio)
print('%d by %d, w/h' % (width, height))

data = np.zeros((height, width, 3))
MAX_DEPTH = 1
DIFFUSE_RAYS = 4
DIFFUSE_DEPTH = 1
DIFFUSE_DEPTH_OFFSET = MAX_DEPTH - DIFFUSE_DEPTH

FOV = 70.0

fov = FOV * math.pi / 180.0

upscale_factor = 2

# gamma
random_light_spread = 0.6
G = 2.2
brightness = 0.1
contrast = 0.5
EV = 150

# alpha for phong
alpha = 10.0
bias = 0.0
vert_division = 10
hor_division = 10


class Ray:
    ray_direction = np.array((0.0, 0.0, -1.0))
    ray_origin = np.array((0.0, 0.0, 0.0))

    refraction_index = 1.0

    def __init__(self, _dir, _or, _ior=None):
        self.ray_direction = _dir
        self.ray_origin = _or
        if _ior is not None:
            self.refraction_index = _ior

    def reflect(self, normal):
        return self.ray_direction - 2 * normal * np.inner(self.ray_direction, normal)

    def point_at_distance(self, _dist):
        return self.ray_origin + _dist * self.ray_direction

    def refract(self, ior, normal):
        Nrefr = normal
        n_dot_incident = np.inner(Nrefr, self.ray_direction)
        ior_1 = self.refraction_index
        ior_2 = ior
        if n_dot_incident < 0:
            # outside
            n_dot_incident *= -1
        else:
            # we are inside, swap normal
            Nrefr = -normal
            ior_1 = ior
            ior_2 = self.refraction_index

        mu = ior_2 / ior_1

        k = 1 - mu ** 2 * (1 - n_dot_incident ** 2)

        if k < 0:
            k *= -1
        return mu * self.ray_direction + mu * (n_dot_incident - math.sqrt(k)) * Nrefr


class Object:
    color = np.array((1.0, 1.0, 1.0))
    center = np.array((0, 0, 0))

    radius = 0

    spec = 0.5
    emissivity = 0
    lambert = 0.6
    ambient = 0.18
    refraction_index = 0.0
    alpha = 1.0
    max_int_dist = 0.0

    type = 'object'

    def __init__(self, color=None):
        if color is not None:
            self.color = color
        self.max_internal_distance()

    def intersect(self, ray):
        print('super class has no intersect defined')

    def normal_at_point(self, point):
        print('super class has no normal defined')

    def intersection_point(self, ray):
        print('super class has no intersection_point defined')

    def max_internal_distance(self):
        print('super class has no max_internal_distance func')


class Sphere(Object):
    radius = 1.0
    type = 'sphere'

    emissivity = 0.0

    def __init__(self, color=None, radius=None):
        if radius is not None:
            self.radius = radius
        super(Sphere, self).__init__(color)

    def intersect(self, ray):
        part_d = np.dot(ray.ray_direction, np.subtract(ray.ray_origin, self.center)) ** 2 - (
                np.linalg.norm(np.subtract(ray.ray_origin, self.center)) ** 2 - self.radius ** 2)
        if part_d >= 0:
            return self.distance(ray, part_d)

        else:
            return math.inf

    def distance(self, ray, part_d):
        d1 = -1.0 * (np.matmul(ray.ray_direction, np.subtract(ray.ray_origin, self.center))) - math.sqrt(part_d)
        # d2 = -1.0*(np.matmul(ray.ray_direction, np.subtract(ray.ray_origin, self.center))) + math.sqrt(part_d)
        # d2 will be larger dan d1
        return d1

    def intersection_point(self, ray):
        t = self.intersect(ray)
        the_point = np.add(ray.ray_origin, ray.ray_direction * t)
        return the_point

    def normal_at_point(self, point):
        return normalise(np.subtract(point, self.center))

    def max_internal_distance(self):
        self.max_int_dist = self.radius


class Plane(Object):
    normal = np.array((0.0, 1.0, 0.0))
    type = 'plane'

    def __init__(self, color=None, normal=None):
        if normal is not None:
            self.normal = normal
        super(Plane, self).__init__(color)

    def intersect(self, ray):
        dot_product = np.inner(self.normal, ray.ray_direction)
        if dot_product == 0:
            return math.inf
        else:
            return self.distance(ray)

    def distance(self, ray):
        distance = np.linalg.norm(ray.ray_origin - self.intersection_point(ray))
        return distance

    def intersection_point(self, ray):
        _N = np.inner(self.normal, ray.ray_direction)
        d = 0
        if _N == 0:
            d = math.inf
        else:
            d = np.inner((self.center - ray.ray_origin), self.normal) / np.inner(self.normal, ray.ray_direction)
        if d > 0:
            return ray.point_at_distance(d)
        else:
            return np.array((math.inf, math.inf, math.inf))

    def normal_at_point(self, point):
        return self.normal

    def max_internal_distance(self):
        self.max_int_dist = math.inf


class Light:
    direction = np.array((-1.0, -1.0, 0.0))
    color = np.array((1.0, 1.0, 1.0))
    origin = np.array((-1.0, -1.0, 0))
    intensity = 1.0

    def __init__(self, direction=None):
        if direction is not None:
            self.direction = direction

    def set_origin(self, x, y, z):
        self.origin = np.array((x, y, z))

    def set_direction(self, x, y, z):
        self.direction = np.array((x, y, z))


def calc_normal(_p1, _p2, _p3):
    P1P2 = _p2 - _p1
    P1P3 = _p3 - _p1
    return normalise(np.cross(P1P2, P1P3))


class Triangle(Plane):
    p2 = np.array((1.0, 0.0, -1.0))
    p3 = np.array((1.0, 1.0, -1.0))

    bary_center = np.array((0.0, 0.0, 0.0))

    A = 0.0

    type = 'triangle'

    def __init__(self, color, normal, _p1, _p2, _p3):
        if _p1 is not None:
            self.center = _p1
        if _p2 is not None:
            self.p2 = _p2
        if _p3 is not None:
            self.p3 = _p3
        self.A = self.area(_p1, _p2, _p3)
        self.bary_center = (_p1 + _p2 + _p3) / 3.0
        super(Triangle, self).__init__(color, normal)

    def intersect(self, ray):
        int_pt = self.intersection_point(ray)

        if np.linalg.norm(int_pt) < math.inf:
            A1 = self.area(int_pt, self.center, self.p2)
            A2 = self.area(int_pt, self.p2, self.p3)
            A3 = self.area(int_pt, self.center, self.p3)

            if abs(A1 + A2 + A3 - self.A) < 0.0001:
                return self.distance(ray)
            else:
                return math.inf
        else:
            return math.inf

    def area(self, P1, P2, P3):
        P1P2 = P2 - P1
        P1P3 = P3 - P1
        return np.linalg.norm(np.cross(P1P2, P1P3)) * 0.5

    def max_internal_distance(self):
        d_p1p2 = np.linalg.norm(self.center - self.p2)
        d_p2p3 = np.linalg.norm(self.p2 - self.p3)
        d_p3p1 = np.linalg.norm(self.p3 - self.center)
        self.max_int_dist = np.amax((d_p1p2, d_p2p3, d_p3p1))


def normalise(vector):
    norm = np.linalg.norm(vector)
    if norm != 0:
        return vector * (1 / np.linalg.norm(vector))


s1 = Sphere(np.array((0.65, 0.44, 0.39)), 1.0)
s1.center = np.array((-2.0, 1.0, -5.0))
s1.lambert = 0.2
s1.spec = 0.0
s1.ambient = 0.18
s1.refraction_index = 1.4
s1.alpha = 1.0
s1.emissivity = 1.0

s2 = Sphere(np.array((1.0, 1.0, 1.0)), 1.0)
s2.center = np.array((0.0, 1.0, -5.0))
s2.lambert = 1.0
s2.ambient = 0.18
s2.spec = 0.0
s2.alpha = 1.0
s2.refraction_index = 1.4
s2.emissivity = 1.0

s3 = Sphere(np.array((1.0, 0.78, 0.34)), 1.0)
s3.center = np.array((2.0, 1.0, -3.0))
s3.lambert = 0.2
s3.ambient = 0.18
s3.spec = 0.5
s3.alpha = 1.0
s3.refraction_index = 1.4
s3.emissivity = 1.0

s4 = Sphere(np.array((0.0, 1.0, 1.0)), 0.4)
s4.center = np.array((4.0, 3.0, -2.0))
s4.lambert = 0.0
s4.ambient = 0.18
s4.spec = 1.0
s4.alpha = 0.0
s4.refraction_index = 1.4
s4.emissivity = 1.0
s4.type = 'light'

s5 = Sphere(np.array((0.95, 0.5, 0.2)), 0.4)
s5.center = np.array((4.0, 3.0, -6.0))
s5.lambert = 0.0
s5.ambient = 0.18
s5.spec = 1.0
s5.alpha = 0.0
s5.refraction_index = 1.4
s5.emissivity = 1.0
s5.type = 'light'


# plane
p1 = Plane(np.array((1.0, 1.0, 1.0)), np.array((0.0, 1.0, 0.0)))
p1.center = np.array((0.0, 0.0, 0.0))
p1.spec = 0.0
p1.lambert = 1.0
p1.ambient = 0.18
p1.alpha = 1.0
p1.refraction_index = 1.4

# triangle points
spiegel = -10
t_pt1 = np.array((-5.0, 0.0, spiegel))
t_pt2 = np.array((5.0, 0.0, spiegel))
t_pt3 = np.array((0.0, 5.0, spiegel))
test_triangle_color = np.array((1.0, 1.0, 1.0))

test_triangle = Triangle(test_triangle_color, calc_normal(t_pt1, t_pt2, t_pt3), t_pt1, t_pt2, t_pt3)
test_triangle.lambert = 0.0
test_triangle.ambient = 0.0
test_triangle.spec = 1.0
test_triangle.emissivity = 1.0

# triangle 2
tt_pt1 = np.array((-5.0, 0.0, -spiegel))
tt_pt2 = np.array((5.0, 0.0, -spiegel))
tt_pt3 = np.array((0.0, 5.0, -spiegel))
test_triangle2_color = np.array((1.0, 1.0, 1.0))

test_triangle2 = Triangle(test_triangle_color, calc_normal(tt_pt1, tt_pt2, tt_pt3), tt_pt1, tt_pt2, tt_pt3)
test_triangle2.lambert = 0.1
test_triangle2.ambient = 0.1
test_triangle2.spec = 1.0

# triangle 3
et_pt1 = np.array((4.0, 0.0, -4.0))
et_pt2 = np.array((4.0, 2.0, -3.0))
et_pt3 = np.array((4.2, 2.0, -6.0))
test_triangle_color = np.array((1.0, 1.0, 1.0))
test_triangle3 = Triangle(test_triangle_color, calc_normal(et_pt1, et_pt2, et_pt3), et_pt1, et_pt2, et_pt3)
test_triangle3.lambert = 0.1
test_triangle3.ambient = 0.0
test_triangle3.spec = 1.0
test_triangle3.emissivity = 1.0
test_triangle3.refraction_index = 1.4
test_triangle3.type = 'light'

# camera
camera_point = np.array((0.0, 3.0, 5.0))
view_direction = normalise(np.array([0.0, -0.1, -1.0]))

scene = [s1, s2, s3, p1, test_triangle, s4, test_triangle3, s5]

light1 = Light(np.array((0.0, -1.0, 0.0)))
light1.set_origin(1.0, 2.3, -3.0)
light1.intensity = 0.0
light1.color = np.array((1.0, 1.0, 1.0))

lights = []
lights.append(light1)

surface_lights = [s4, s5,test_triangle3]



# for x in range(0, 10):
#    random_loc = np.random.rand(3) * 2.0 + np.array((0, 0, 1))
#    random_radius = np.random.rand(1) * 2.0
#    random_color = np.random.rand(3)
#
#    random_sphere = Sphere(random_color, random_radius)
#    random_sphere.center = random_loc
#
#    # random specs
#    random_specs = np.random.rand(3)
#    random_sphere.lambert = random_specs[0]
#    random_sphere.spec = random_specs[1]
#    random_sphere.ambient = random_specs[2]
#
#    scene.append(random_sphere)


def make_rot_matrix(rx, ry, rz):
    cx, sx = np.cos(rx), np.sin(rx)
    cy, sy = np.cos(ry), np.sin(ry)
    cz, sz = np.cos(rz), np.sin(rz)
    Ry = np.array(((cy, 0, sy), (0, 1, 0), (-sy, 0, cy)))
    Rx = np.array(((1, 0, 0), (0, cx, -sx), (0, sx, cx)))
    Rz = np.array(((cz, -sz, 0), (sz, cz, 0), (0, 0, 1)))
    return np.dot(Rz, np.dot(Ry, Rx))


def rotate_vector(vect, rot_mat):
    return np.dot(rot_mat, vect)


def make_rot_matrix_from_normals(_u, _v):
    c = np.inner(_u, _v)
    _R = np.zeros((3, 3)) - np.identity(3)
    if c == 1:
        _R = np.identity(3)
    if abs(c) != 1:
        _a = np.cross(_u, _v)
        _a_cross = np.array(((0, -_a[2], _a[1]),
                             (_a[2], 0, -_a[0]),
                             (-_a[1], _a[0], 0)))
        _a_cross_sq = np.dot(_a_cross, _a_cross)
        _R = np.identity(3) + _a_cross + _a_cross_sq * 1.0 / (1 + c)
    return _R


def interpolated_ray_dir(_fov, _y, _x, _width, _height, view_dir):
    y_clipped = -(_y / _height - 0.5)
    x_clipped = (_x / _width - 0.5) * aspect_ratio

    D = 1.0 / math.tan(_fov / 2.0)
    uv = (np.array((x_clipped, y_clipped, 0)))

    return normalise(np.add(uv, D * view_dir))


def opening_corner(_n, _d, _L, _dist):
    _gamma = math.pi * 0.5 - abs(np.inner(_n, _d))
    _delta = math.atan(math.sin(_gamma) / (_dist / _L - math.cos(_gamma)))
    return (1.0 - _delta / (math.pi) * 2)


def intersect(_ray, _scene):
    closest = math.inf
    _nearest_object = None
    for c in _scene:
        distance_to_c = c.intersect(_ray)
        if distance_to_c < closest:
            if distance_to_c > 0.001:
                closest = distance_to_c
                _nearest_object = c
    return [_nearest_object, closest]


def diffuse_cos_weighted_direction(_spec, _reflection_vector, _normal):
    m = alpha ** (_spec)
    _u = random.random()
    _v = random.random()

    sinT = math.sqrt(1.0 - ((1.0 - _u) ** (1.0 / (1 + m)) ** 2))
    phi = 2 * math.pi * _v

    _X = sinT * math.cos(phi)
    _Y = sinT * math.sin(phi)
    _Z = _u

    random_vector = np.array((_X, _Y, _Z))
    # rotate so it is on the normal plane
    _rotation_mat = make_rot_matrix_from_normals(np.array((0, 0, 1.0)), _normal)
    _random_rotated_vector = rotate_vector(random_vector, _rotation_mat).reshape(_reflection_vector.shape)

    # weigh it with the _reflectionvector
    weighted = normalise(_spec * _reflection_vector + (1 - _spec) * _random_rotated_vector)
    return weighted


def trace(_ray, _scene, _depth):
    col = np.array((0.0, 0.0, 0.0))
    if _depth <= MAX_DEPTH:
        [intersect_object, distance] = intersect(_ray, _scene)
        if intersect_object is not None:

            if intersect_object.type == 'light':
                return intersect_object.color * intersect_object.emissivity

            intersection_point = _ray.point_at_distance(distance)
            surf_norm = intersect_object.normal_at_point(intersection_point)

            # direct lighting
            direct_col = np.array((0.0, 0.0, 0.0))
            for s in surface_lights:
                for o in range(0, DIFFUSE_RAYS):
                    pt_to_light_vec0 = normalise(s.center - intersection_point)
                    random_direction0 = diffuse_cos_weighted_direction(random_light_spread, pt_to_light_vec0,
                                                                       pt_to_light_vec0)
                    pt_to_light_ray0 = Ray(random_direction0, intersection_point)

                    [towards_light_intersect0, distance_pt_to_light0] = intersect(pt_to_light_ray0, _scene)
                    if towards_light_intersect0 is not None:
                        if towards_light_intersect0.type == 'light':
                            n_dot_pt_l = np.inner(surf_norm, pt_to_light_vec0) * intersect_object.lambert
                            if n_dot_pt_l > 0:
                                distance_fall_off = 4 * math.pi * distance_pt_to_light0 ** 2
                                direct_col += intersect_object.color * trace(pt_to_light_ray0, _scene,
                                                                             MAX_DEPTH) * n_dot_pt_l / distance_fall_off
            if DIFFUSE_RAYS != 0:
                direct_col /= DIFFUSE_RAYS
                col += direct_col
            # amb

            # lambertion point light
            for li in lights:
                if li.intensity > 0:
                    pt_to_light_dist = np.linalg.norm(light1.origin - intersection_point)
                    pt_to_light_vec = normalise(light1.origin - intersection_point)
                    pt_to_light_ray = Ray(pt_to_light_vec, intersection_point)

                    [towards_light_intersect, distance_pt_to_light] = intersect(pt_to_light_ray, _scene)

                    if distance_pt_to_light is math.inf:  # it sees the light
                        lamb_int = np.inner(surf_norm, pt_to_light_vec)
                        if lamb_int > 0:
                            col += intersect_object.color * light1.intensity * light1.color / (
                                    4.0 * math.pi * pt_to_light_dist ** 2) * intersect_object.lambert * lamb_int

            # spec
            reflected_ray = Ray(normalise(_ray.reflect(surf_norm)), intersection_point)
            dice = random.random()
            if dice < intersect_object.spec:
                random_dir = diffuse_cos_weighted_direction(intersect_object.spec, reflected_ray.ray_direction,
                                                            reflected_ray.ray_direction)
                col += trace(Ray(random_dir, reflected_ray.ray_origin),
                             _scene, _depth + 1) * intersect_object.color / intersect_object.spec * np.inner(
                    random_dir, surf_norm)
            else:
                # random ray
                if DIFFUSE_RAYS != 0:
                    diffuse_color = np.array((0.0, 0.0, 0.0))
                    for i in range(0, DIFFUSE_RAYS):
                        random_direction = diffuse_cos_weighted_direction(0.0,
                                                                          reflected_ray.ray_direction,
                                                                          surf_norm)
                        diffuse_color += intersect_object.color * trace(Ray(random_direction, intersection_point),
                                                                        _scene,
                                                                        _depth + 1 + DIFFUSE_DEPTH_OFFSET) / (
                                                 1 - intersect_object.spec) * np.inner(random_direction, surf_norm)

                    col += diffuse_color / float(DIFFUSE_RAYS)

            # refraction
            if intersect_object.alpha != 1.0:
                new_ior = 1.0
                if _ray.refraction_index == 1.0:
                    new_ior = intersect_object.refraction_index
                else:
                    surf_norm *= -1
                refracted_ray = Ray(normalise(_ray.refract(new_ior, surf_norm)),
                                    intersection_point + bias * _ray.ray_direction, new_ior)
                col += trace(refracted_ray, _scene, _depth + 1)

            return col
        else:
            return col
    else:
        return col


view_direction = normalise(view_direction)


def render_pixel(_fov, _x, _y, _width, _height, _view_direction, _viewpoint, _scene):
    return trace(Ray(interpolated_ray_dir(fov, _x, _y, _width, _height, _view_direction), _viewpoint), _scene, 0)


class Render_thread(threading.Thread):
    x_start = 0
    x_end = height
    y_start = 0
    y_end = width
    m_scene = []

    render_data = np.zeros((x_end, y_end, 3))

    def __init__(self, threadID, name, x_s=None, x_e=None, y_s=None, y_e=None, _scene=None):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.name = name
        self.x_start = x_s
        self.x_end = x_e
        self.y_end = y_e
        self.y_start = y_s
        self.m_scene = _scene
        self.render_data = np.zeros((x_e - x_s, y_e - y_s, 3))

        print('making a thread with hor from %.f to %.f and height from %.f to %.f' % (
            self.y_start, self.y_end, self.x_start, self.x_end))

    def run(self):
        for _i in range(self.x_start, self.x_end):
            for _k in range(self.y_start, self.y_end):
                self.render_data[_i - self.x_start, _k - self.y_start,] = render_pixel(fov, _i, _k, width, height,
                                                                                       view_direction, camera_point,
                                                                                       self.m_scene)
        self.render_data **= (1 / G)


# threads = []
# num_of_threads = vert_division * hor_division
# delta_x = math.floor(height / vert_division)
# delta_y = math.floor(width / hor_division)
# for k in range(0, vert_division):
#    for j in range(0, hor_division):
#        h_st = delta_x * k
#        w_st = delta_y * j
#        h_end = h_st + delta_x
#        w_end = w_st + delta_y
#        thread_name = 'Thread-' + str(k + j)
#        threads.append(Render_thread(k + j, thread_name, h_st, h_end, w_st, w_end, scene))
#
# t_start = time.time()
# for idx in range(len(threads)):
#    threads[idx].start()
# for idx in range(len(threads)):
#    threads[idx].join()
# t_end = time.time()
# print('image rendered with %d threads in %f s with width = %d and height = %d' % (
# num_of_threads, (t_end - t_start), width, height))
#
# stitched_data = np.zeros((height, width, 3))
#
#
# def stitch_buffers(_threads):
#    for thr in _threads:
#        x_starts = thr.x_start
#        y_starts = thr.y_start
#        x_ends = thr.x_end
#        y_ends = thr.y_end
#        for _i in range(x_starts, x_ends):
#            for _j in range(y_starts, y_ends):
#                stitched_data[_i, _j,] = thr.render_data[_i-x_starts, _j-y_starts,]
#
#
# stitch_buffers(threads)
# stitched_data *= 255 / np.amax(stitched_data)
# img = Image.fromarray(stitched_data.astype(np.uint8))
# img.show()

print('Rendering')
t_start = time.time()
for x in range(0, height):
    for y in range(0, width):
        data[x, y,] = render_pixel(fov, x, y, width, height, view_direction, camera_point, scene)
t_end = time.time()

print('image rendered in %f s with width = %d and height = %d' % ((t_end - t_start), width, height))

# gamma correct


data /= np.amax(data)
data *= EV
data = data * (1.0 + data / EV) / (data + 1.0)
data = np.clip(data, 0, 1)
data **= (1 / G)
data *= 255
img = Image.fromarray(data.astype(np.uint8))
print('upscaling')
t_start = time.time()
img2 = img.resize([upscale_factor * width, upscale_factor * height], Image.LANCZOS)
t_end = time.time()
print('image upscaled in %f s with  factor %d' % ((t_end - t_start), upscale_factor))
img2.show()
img2.save('image.png')