#!/usr/bin/python3

from math import pi, cos, sin, exp

import phue, pygame, colorsys

# === graph data ===

# vertices as screen positions
viewsize = 600
shrink = 0.9
vertices = [(viewsize*(1 + shrink*x)/2, viewsize*(1 + shrink*y)/2) for x, y in
  [(0, 0)] + [(sin(n*pi/5), -cos(n*pi/5)) for n in range(10)] + [(0, 0)]
]

# edges as vertex pairs. FireStar will rest on vertices 6 (A), 5 (B), 7 (C)
vertex_adj = [
  (0, 1), (0, 3), (0, 5), (0, 7), (0, 9),
  (2, 8), (2, 6), (6, 10), (10, 4), (4, 8),
  (2, 9), (2, 11), (2, 5),
  (4, 1), (4, 11), (4, 7),
  (6, 3), (6, 11), (6, 9),
  (8, 5), (8, 11), (8, 1),
  (10, 7), (10, 11), (10, 3),
  (9, 3), (9, 5),
  (1, 5), (1, 7),
  (3, 7)
]

# cyclic edge adjacencies
cyc_edge_adj = [
  (0, 1), (1, 2), (2, 3), (3, 4), (4, 0),
  (5, 10), (10, 11), (11, 12), (12, 6), (6, 5),
  (6, 16), (16, 17), (17, 18), (18, 7), (7, 6),
  (7, 22), (22, 23), (23, 24), (24, 8), (8, 7),
  (8, 13), (13, 14), (14, 15), (15, 9), (9, 8),
  (5, 9), (9, 19), (19, 20), (20, 21), (21, 5),
  (4, 10), (10, 25), (25, 26), (26, 18), (18, 4),
  (3, 15), (15, 29), (29, 28), (28, 22), (22, 3),
  (2, 12), (12, 27), (27, 26), (26, 19), (19, 2),
  (1, 24), (24, 25), (25, 29), (29, 16), (16, 1),
  (0, 21), (21, 28), (28, 27), (27, 13), (13, 0),
  (14, 11), (11, 23), (23, 20), (20, 17), (17, 14)
]

# === pygame setup ===

pygame.init()
screen = pygame.display.set_mode((viewsize, viewsize + 100))
pygame.display.set_caption('Preview')
font = pygame.font.Font(None, 25)
clock = pygame.time.Clock()

# === bridge setup and control ===

# --- copied from firestar.py

b = phue.Bridge('172.18.130.12')
##b = None
Nedges = 30;
edgecode = [None] * Nedges;

# The edges are numbered as in info/SSDedges.obj; enter the two-letter code for
# the lightstrip inserted through each edge in the block below:

exec(open("edgecodes.py").read())

if b != None:
  edges = [ b[code] for code in edgecode ]

lasthue = 65535
fullbright = 254
fullsat = 254

# --- control methods and diagnostics

updates = 0
new_updates = 0

def set_light_hsv(e, color, hit_bridge):
  command = {
    'hue' : int(color[0]*lasthue),
    'bri' : int(color[2]*fullbright),
    'sat' : int(color[1]*fullsat)
  }
  if hit_bridge:
    if b != None:
      b.set_light(edges[e].light_id, command)
    global new_updates
    new_updates += 1
  rgb_color = tuple(255*c for c in colorsys.hsv_to_rgb(*color))
  pygame.draw.aaline(screen, rgb_color, vertices[vertex_adj[e][0]], vertices[vertex_adj[e][1]])

# === messages and mode selection ===

mode = 0

# data labels
updates_label = font.render('updates', True, (255, 255, 255))
energy_label = font.render('energy', True, (255, 255, 255))
position_label = [
  font.render('temp[0]', True, (255, 255, 255)),
  font.render('charge[0]', True, (255, 255, 255))
]

# instructions
mode_inds = [
  [
    font.render('[d] heat', True, (255, 255, 255)),
    font.render('[f] wave', True, (150, 150, 150))
  ],
  [
    font.render('[d] heat', True, (150, 150, 150)),
    font.render('[f] wave', True, (255, 255, 255))
  ]
]
bop_instr = [
  font.render('press [j] to bop', True, (255, 255, 255)),
  font.render('press [j] or [k] to bop', True, (255, 255, 255))
]
drain_instr = font.render('hold [space] to drain', True, (255, 255, 255))

# set up background
background = pygame.Surface(screen.get_size()).convert()
background.fill((32, 32, 32))
background.blit(drain_instr, (340, 60 + viewsize))

# === interaction parameters ===

# these frame rate and threshold settings are untested, but the exponential
# moving average update frequency suggests they should be just barely in bounds
frame_rate = 10
ema_weight = 2/(frame_rate+1)

inductance = 0.001

therm_conduct = 0.01

drain = False
drain_rate = 0.1

# === phase space ===

# heat equation
temp = 30*[0]
new_temp = 30*[0]

# wave equation
charge = 12*[0]
current = 30*[0]

litness = 30*[float("inf")]
change_threshold = 0.15

# === main loop ===

def heat_evolution():
  for e, f in cyc_edge_adj:
    flow = therm_conduct*(temp[e] - temp[f])
    new_temp[e] -= flow
    new_temp[f] += flow
  for e in range(30):
    if drain:
      new_temp[e] *= 1 - drain_rate
    temp[e] = max(new_temp[e], 0)

def wave_evolution():
  # use verlet integration, first updating the currents and then using the new
  # currents to update the charges.
  for e in range(30):
    u, v = vertex_adj[e]
    current[e] += inductance*(charge[u] - charge[v])
  for e in range(30):
    u, v = vertex_adj[e]
    charge[u] -= current[e]
    charge[v] += current[e]
  if drain:
    for v in range(12):
      charge[v] *= 1 - drain_rate

def set_heat_light(e, i):
  new_litness = 1 - exp(-2*i)
  if abs(new_litness - litness[e]) > change_threshold:
    hit_bridge = True
    litness[e] = new_litness
  else:
    hit_bridge = False
  set_light_hsv(e, (0.167*litness[e], 1 - 0.5*litness[e]*litness[e], litness[e]), hit_bridge)

def set_wave_light(e, i):
  new_litness = 1 - exp(-2*i*i)
  if abs(new_litness - litness[e]) > change_threshold:
    hit_bridge = True
    litness[e] = new_litness
  else:
    hit_bridge = False
  set_light_hsv(e, (0.45 + litness[e]*0.25, 1 - 0.5*litness[e]*litness[e], litness[e]), hit_bridge)

def handle_heat_events():
  global drain, mode
  for event in pygame.event.get():
    if event.type == pygame.KEYDOWN:
      if event.key == pygame.K_j:
        new_temp[0] += 10
      elif event.key == pygame.K_SPACE:
        drain = True
      elif event.key == pygame.K_f:
        mode = 1
    elif event.type == pygame.KEYUP:
      if event.key == pygame.K_SPACE:
        drain = False
    elif event.type == pygame.QUIT:
      quit()

def handle_wave_events():
  global drain, mode
  for event in pygame.event.get():
    if event.type == pygame.KEYDOWN:
      if event.key == pygame.K_j:
        charge[0] += 100
        charge[11] -= 100
      elif event.key == pygame.K_k:
        charge[5] += 100
        charge[7] -= 100
      elif event.key == pygame.K_SPACE:
        drain = True
      elif event.key == pygame.K_d:
        mode = 0
    elif event.type == pygame.KEYUP:
      if event.key == pygame.K_SPACE:
        drain = False
    elif event.type == pygame.QUIT:
      quit()

def show_data(screen):
  # light updates
  global updates, new_updates
  updates = ema_weight*new_updates + (1 - ema_weight)*updates
  new_updates = 0
  updates_value = font.render(str(frame_rate * updates), True, (255, 255, 255))
  screen.blit(updates_label, (20, viewsize))
  screen.blit(updates_value, (120, viewsize))
  
  # energy
  if mode == 0:
    energy = sum(temp)
  elif mode == 1:
    energy = 0.5*(sum(q*q for q in charge) + inductance*sum(i*i for i in current))
  energy_value = font.render(str(energy), True, (255, 255, 255))
  screen.blit(energy_label, (20, 30 + viewsize))
  screen.blit(energy_value, (120, 30 + viewsize))
  
  # a position
  ##position_value = None
  if mode == 0:
    position_value = font.render(str(temp[0]), True, (255, 255, 255))
  elif mode == 1:
    position_value = font.render(str(charge[0]), True, (255, 255, 255))
  screen.blit(position_label[mode], (20, 60 + viewsize))
  screen.blit(position_value, (120, 60 + viewsize))
  
  # instructions
  screen.blit(mode_inds[mode][0], (340, viewsize))
  screen.blit(mode_inds[mode][1], (420, viewsize))
  screen.blit(bop_instr[mode], (340, 30 + viewsize))

if __name__ == '__main__':
  while True:
    # handle events
    if mode == 0:
      handle_heat_events()
    elif mode == 1:
      handle_wave_events()
    
    # clear screen
    screen.blit(background, (0, 0))
    
    # show data and instructions
    show_data(screen)
    
    # show state and evolve state
    if mode == 0:
      for e in range(30):
        set_heat_light(e, temp[e])
      heat_evolution()
    elif mode == 1:
      for e in range(30):
        set_wave_light(e, current[e])
      wave_evolution()
    
    # step
    pygame.display.flip()
    clock.tick(frame_rate)