from PIL import Image, ImageDraw
import argparse
import sys
import math

def getNodePosition(i, nodes, deforms):
  (x, y) = nodes[i]
  return (
    x + deforms[2 * i + 0] * 50.0,
    y + deforms[2 * i + 1] * 50.0
  )

def PostProcess(inputfile, outputfile):
  content = []
  with open(inputfile) as f:
    content = f.readlines()

  line = 1
  nodesCount = int(content[line])
  line+=1

  nodes = []

  for i in range(0, nodesCount):
    coordinates = content[line].split()
    line+=1
    nodes.append(
      (
        float(coordinates[0]),
        float(coordinates[1])
      )
    )

  (max_x, max_y) = nodes[0]
  (min_x, min_y) = nodes[0]

  for node in nodes:
    (x, y) = node
    if max_x < x:
      max_x = x
    if min_x > x:
      min_x = x
    if max_y < y:
      max_y = y
    if min_y > y:
      min_y = y

  elementsCount = int(content[line])
  line+=1

  elements = []

  for i in range(0, elementsCount):
    elementNodes = content[line].split()
    line+=1
    elements.append(
      (
        int(elementNodes[0]),
        int(elementNodes[1]),
        int(elementNodes[2])
      )
    )

  center = [(max_x + min_x) / 2.0, (max_y + min_y) / 2.0]
  scale = max_x - min_x;
  if scale < max_y - min_y:
    scale = max_y - min_y

  deforms = []
  stresses = []
  resultsContent = []
  with open(outputfile) as f:
    resultsContent = f.readlines()

  for line in resultsContent[:nodesCount*2]:
    deforms.append(float(line))
  for line in resultsContent[nodesCount*2:]:
    stresses.append(float(line))

  max_stress = 0
  min_stress = stresses[0]
  for stress in stresses:
    if max_stress < stress:
      max_stress = stress
    if min_stress > stress:
      min_stress = stress
  print("Max stress: " + str(max_stress))

  def Transform(point):
    x = (point[0] - center[0]) / scale / 1.5 * image.size[0] + image.size[0]/2.0
    y = (point[1] - center[1]) / scale / 1.5 * image.size[1] + image.size[1]/2.0
    return (x, image.size[1] - y)

  image = Image.new("RGB", (1024,1024))

  draw = ImageDraw.Draw(image)
  draw.rectangle([(0, 0), image.size], fill=0x777777)

  for element in elements:
    draw.polygon([Transform(nodes[element[0]]), Transform(nodes[element[1]]), Transform(nodes[element[2]])], fill=0x20a020, outline=0x707020)

  image.save("initial.png", "PNG")

  draw.rectangle([(0, 0), image.size], fill=0x777777)

  rangesIntervals = 12

  if max_stress == min_stress:
    max_stress += 1.0
    min_stress -= 1.0

  for idx,element in enumerate(elements):
    v = (stresses[idx] - min_stress) / (max_stress - min_stress)
    #v = math.floor(v * (rangesIntervals + 1.0))/rangesIntervals;
    colorf = ((v - 0.5)*4.0, 2.0 - abs(v- 0.5)*4.0,  2.0 - v*4.0)
    colori = (int(colorf[0]*255), int(colorf[1]*255), int(colorf[2]*255))
    draw.polygon(
      [
        Transform(getNodePosition(element[0], nodes, deforms)),
        Transform(getNodePosition(element[1], nodes, deforms)),
        Transform(getNodePosition(element[2], nodes, deforms))
      ],
      fill = colori
    )

  image.save("deformed.png", "PNG")


if __name__ == '__main__':
  parser = argparse.ArgumentParser(description='Tool postprocessing output of MinimalFem.',
                   epilog = 'Visit http://podgorskiy.com/spblog/392/writing-a-fem-solver-in-less-the-180-lines-of-code')
  parser.add_argument('inputfile', metavar='inputfile', type=str, help='Path to input file.')
  parser.add_argument('outputfile', metavar='outputfile', type=str, help='Path to output file.')
  args = parser.parse_args()
  PostProcess(args.inputfile, args.outputfile)
