# cording: utf-8
# --+----1----+----2----+----3----+----4----+----5----+----6----+----7----+----8
# NAME       : pyrender.py
# AUTHOR     : CASPAR003 <caspar003@caspar003.info>
# DATE       : 2013-12-05
# DESCRIPTION: smallptのPython実装．ぜんぜんsmallではない．
# --+----1----+----2----+----3----+----4----+----5----+----6----+----7----+----8
import math
import numpy as np
from PIL import Image
import random
import sys

RFDIFF = 0
RFSPEC = 1
RFREFR = 2

def main(ns):
  width   = 512
  height  = 512
  gamma   = 2.2
  nsample = ns
  
  cam = ray(np.array([ 50.0, 52.0,295.6]), norm(np.array([0.0,-0.042612,-1.0])))
  cx  = np.array([float(width) * 0.5135 / height, 0.0, 0.0])
  cy  = norm(cross(cx, cam.d)) * 0.5135
  r   = np.zeros(3)
  img = Image.new("RGB", (width, height))
  pix = img.load()
  cr  = np.zeros((width, height))
  cg  = np.zeros((width, height))
  cb  = np.zeros((width, height))
  for y in range(height):    # 行ループ
    print("\rRendering ({0:d} spp) {1:6.2f}%".format(nsample*4, 100.0 * float(y) / float(height - 1)), end="")
    for x in range(width):   # 列ループ
      for sy in range(2):    # サブピクセル行ループ
        for sx in range(2):  # サブピクセル列ループ
          r = np.zeros(3)
          for s in range(nsample):
            r1 = 2.0 * random.random()
            if (r1 < 1.0):
              dx = math.sqrt(r1) - 1.0
            else:
              dx = 1.0 - math.sqrt(2.0 - r1)
            r2 = 2.0 * random.random()
            if (r2 < 1.0):
              dy = math.sqrt(r2) - 1.0
            else:
              dy = 1.0 - math.sqrt(2.0 - r2)
            d = (cx*(((sx + 0.5 + dx)/2.0 + x)/width  - 0.5) +
                 cy*(((sy + 0.5 + dy)/2.0 + y)/height - 0.5) + cam.d)
            r = r + radiance(ray(cam.o + d*140.0, norm(d)), 0)*(1.0/float(nsample))
          cr[x][y] += 0.25 * clamp(r[0])
          cg[x][y] += 0.25 * clamp(r[1])
          cb[x][y] += 0.25 * clamp(r[2])
      pix[x, height - y - 1] = (int(255.0 * cr[x][y]**(1.0/gamma)), int(255.0 * cg[x][y]**(1.0/gamma)), int(255.0 * cb[x][y]**(1.0/gamma)))
  
  print("")
  img.save("render.png")
  
  return 0

def radiance(r, depth):
  t, id = intersect(r)
  if (id < 0): return np.array([0.0, 0.0, 0.0])  # 交差がなければ黒を返す
  obj = spheres[id]
  
  x = r.o + r.d * t    # 交点
  n = norm(x - obj.p)  # 法線
  if (dot(n, r.d) < 0.0):
    nl = n.copy()
  else:
    nl = n.copy() * (-1.0)
  f = obj.c.copy()
  
  if (f[0] > f[1] and f[0] > f[2]):  # 最大の反射係数を再反射率に使用する
    p = f[0]
  elif (f[1] > f[2]):
    p = f[1]
  else:
    p = f[2]
  
  if (random.random() < p):  # ロシアンルーレット
    f /= p
  else:
    return obj.e
  
  if (obj.refl == RFDIFF):  # Diffuse
    r1  = 2.0 * math.pi * random.random()
    r2  = random.random()
    r2s = math.sqrt(r2)
    w = nl
    if (abs(w[0] > 0.1)):
      u = norm(cross(np.array([0.0, 1.0, 0.0]), w))
    else:
      u = norm(cross(np.array([1.0, 0.0, 0.0]), w))
    v = cross(w, u)
    d = norm(u * math.cos(r1) * r2s + v * math.sin(r1) * r2s + w * math.sqrt(1.0 - r2))
    return obj.e + mult(f, radiance(ray(x, d), depth + 1))
  
  if (obj.refl == RFSPEC):  # Specular
    return obj.e + mult(f, radiance(ray(x, r.d - n * 2.0 * dot(n, r.d)), depth + 1))
  
  if (obj.refl == RFREFR):  # Refraction
    refray = ray(x, r.d - n * 2.0 * dot(n, r.d))  # 反射レイ
    into = dot(n, nl) > 0.0  # 球に対して入射するレイかどうか
    nc = 1.0
    nt = 1.5
    if (into):
      nnt = nc/nt
    else:
      nnt = nt/nc
    ddn = dot(r.d, nl)
    cos2t = 1.0 - nnt**2 * (1.0 - ddn**2)
    if (cos2t < 0.0): # 全反射
      return obj.e + mult(f, radiance(refray, depth + 1))
    
    if (into):
      tdir = norm(r.d * nnt - n * (ddn * nnt + math.sqrt(cos2t)))
    else:
      tdir = norm(r.d * nnt + n * (ddn * nnt + math.sqrt(cos2t)))
    
    a = nt - nc
    b = nt + nc
    R0 = a**2/b**2
    if (into):
      c = 1.0 + ddn
    else:
      c = 1.0 - dot(tdir, n)
    Re = R0 + (1.0 - R0) * c**5
    Tr = 1.0 - Re
    P = 0.25 + 0.5 * Re
    RP = Re/P
    TP = Tr/(1.0-P)
    if (random.random() < P):
      return obj.e + mult(f, radiance(refray, depth +1) * RP)
    else:
      return obj.e + mult(f, radiance(ray(x, tdir), depth +1) * Tr)

def intersect(r):
  t = 1.0e+37
  id = -1
  for i in range(len(spheres)):
    d = spheres[i].intersect(r)
    if (d != 0.0 and d < t):
      t  = d
      id = i
  return t, id

class sphere:
  def __init__(self, radius, position, emission, color, reftype):
    self.rad  = radius
    self.p    = position
    self.e    = emission
    self.c    = color
    self.refl = reftype
  def intersect(self, r):
    op = self.p - r.o
    eps = 1.0e-4
    b = dot(op, r.d)
    det = b*b - dot(op, op) + self.rad*self.rad
    if (det < 0.0):
      return 0.0
    else:
      det = math.sqrt(det)
      t = b - det
      if (t > eps): return t
      t = b + det
      if (t > eps): return t
    return 0.0

class ray:
  def __init__(self, org, dir):
    self.o = org
    self.d = dir

def norm(v):
  return v * (1.0 / math.sqrt(v[0]**2 + v[1]**2 + v[2]**2))

def cross(v1, v2):
  return np.array([
    v1[1]*v2[2] - v1[2]*v2[1],
    v1[2]*v2[0] - v1[0]*v2[2],
    v1[0]*v2[1] - v1[1]*v2[0]
  ])

def dot(v1, v2):
  return v1[0]*v2[0] + v1[1]*v2[1] + v1[2]*v2[2]

def mult(v1, v2):
  return np.array([v1[0]*v2[0], v1[1]*v2[1], v1[2]*v2[2]])

def clamp(v):
  if (v < 0.0): return 0.0
  if (v > 1.0): return 1.0
  return v

if (__name__ == "__main__"):
  nsample = 1
  if (len(sys.argv) > 1): nsample = int(sys.argv[1])
  
  spheres = (
    sphere(  # Left
      1.0e+5,
      np.array([1.0e+5 + 1.0, 40.8, 81.6]),
      np.array([0.0, 0.0, 0.0]),
      np.array([0.75, 0.25, 0.25]),
      RFDIFF
    ),
    sphere(  # Right
      1.0e+5,
      np.array([-1.0e+5 + 99.0, 40.8, 81.6]),
      np.array([0.0, 0.0, 0.0]),
      np.array([0.25, 0.25, 0.75]),
      RFDIFF
    ),
    sphere(  # Back
      1.0e+5,
      np.array([50.0, 40.8, 1.0e+5]),
      np.array([0.0, 0.0, 0.0]),
      np.array([0.75, 0.75, 0.75]),
      RFDIFF
    ),
    sphere(  # Front
      1.0e+5,
      np.array([50.0, 40.8, -1.0e+5 + 170]),
      np.array([0.0, 0.0, 0.0]),
      np.array([0.75, 0.75, 0.75]),
      RFDIFF
    ),
    sphere(  # Bottom
      1.0e+5,
      np.array([50.0, 1.0e+5, 81.6]),
      np.array([0.0, 0.0, 0.0]),
      np.array([0.75, 0.75, 0.75]),
      RFDIFF
    ),
    sphere(  # Top
      1.0e+5,
      np.array([50.0, -1.0e+5 + 81.6, 81.6]),
      np.array([0.0, 0.0, 0.0]),
      np.array([0.75, 0.75, 0.75]),
      RFDIFF
    ),
    sphere(  # Mirror
      16.5,
      np.array([27.0, 16.5, 47.0]),
      np.array([0.0, 0.0, 0.0]),
      np.array([0.95, 0.95, 0.95]),
      RFSPEC
    ),
    sphere(  # Glass
      16.5,
      np.array([73.0, 16.5, 78.0]),
      np.array([0.0, 0.0, 0.0]),
      np.array([0.95, 0.95, 0.95]),
      RFREFR
    ),
    sphere(  # Light
      600.0,
      np.array([50.0, 681.6 - 0.27, 81.6]),
      np.array([12.0, 12.0, 12.0]),
      np.array([0.0, 0.0, 0.0]),
      RFDIFF
    )
  )
  
  sys.exit(main(nsample))