# triangles.py
# Andrew Davison, ad@coe.psu.ac.th, April 2025
'''
  Recursive triangles drawing. An equilateral triangle
  is surrounded by 1/2 sized copies on each side.
'''

import math
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from HC2Mat import *

AXIS_SIZE = 120   # size of window in terms of axis size
SCALE = 1/2
SIZE = 60  # of original side
STEPS = 6   # depth of recursion

colors = [ 'skyblue', 'lightgreen', 'violet', 'salmon']
  # cycle through these for different depths


def drawTriangles(ax, coords, depth):
  if depth > 0:
    col = colors[depth%len(colors)]
    drawPoly(ax, coords, col)
    # rotate smaller triangles so their bases are 
    # parallel with this triangle's sides
    drawTriangles(ax, transformCoords(coords, 60), depth-1)
    drawTriangles(ax, transformCoords(coords, 180), depth-1)
    drawTriangles(ax, transformCoords(coords, -60), depth-1)


def transformCoords(coords, angle):
  tCenter = calcCenter(coords)
  xTip = (coords[2][0] - tCenter[0])
  yTip = (coords[2][1] - tCenter[1])
       # offset from center to tip

  # scale, translate outwards, then rotate,
  # relative to the origin
  m = multMatsList([ 
               transMat(tCenter[0], tCenter[1]),
               rotMat(angle), 
               transMat(2*xTip, 2*yTip),
               scaleMat(SCALE, SCALE),
               transMat(-tCenter[0], -tCenter[1]) ])
  return [applyMat(m, c) for c in coords]


def calcCenter(coords):
  xc = (coords[0][0] + coords[1][0] + coords[2][0])/3
  yc = (coords[0][1] + coords[1][1] + coords[2][1])/3
  return (xc, yc)


def drawPoly(ax, coords, col):
  # colored polygon, edged in black
  poly = Polygon(coords, closed=True, facecolor=col,
                   edgecolor='black', lw=0.5)
  ax.add_patch(poly)


# --------------------------------------

fig, ax = plt.subplots(figsize=(6, 6))
ax.set_aspect('equal')
ax.set_xlim(-AXIS_SIZE, AXIS_SIZE)
ax.set_ylim(-AXIS_SIZE, AXIS_SIZE)
ax.axis('off')

# the initial triangle is equilateral
# with SIZE lengths, centered on the origin
base = math.sqrt(3)/6
coords = [ (-0.5,-base), (0.5,-base), (0, 2*base)]  # unit lengths
coords = [applyMat(scaleMat(SIZE,SIZE), c) for c in coords]

drawTriangles(ax, coords, STEPS)

plt.show()