"""
Sierpinski Gasket - Level-n Construction with Chaos Game Points
===================================================
Draws the level-n Sierpinski gasket (n = 2..5) and plots all chaos game
points from (2+n)-tuples of moves {1,2,3} starting from, for example vertex V_k
that land inside interior of a small triangle.

Barycentric coordinates:
  V_1 = (1, 0, 0)  [top]
  V_2 = (0, 1, 0)  [lower left]
  V_3 = (0, 0, 1)  [lower right]

Chaos game rule: each move goes halfway toward the chosen vertex.
After (2+n) moves starting from V_k, the denominator is 2^(n+2) = D.

  P = V_k/D + sum_{i=1}^{n+2}  V_{M[i]} / 2^(n+3-i)

A point (with coords expressed as integers over D) is in the interior of a 
small triangle iff none of its barycentric coordinates is a multiple
of 1/2^n, i.e iff no numerator is divisible by D/2^n = 4.
"""

# ── imports ─────────────────────────────────────────
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from fractions import Fraction
from itertools import product as iproduct

# ── user input ────────────────────────────────────────
while True:
    try:
        n = int(input("Enter level n (2 to 5): "))
        if 2 <= n <= 5:
            break
        print("  Please enter an integer between 2 and 5.")
    except ValueError:
        print("  Invalid input – please enter an integer.")
while True:
    try:
        initV = int(input("Enter initial vertex (1,2,or 3): "))
        if 1 <= initV <= 3:
            break
        print("  Please enter an integer between 1 and 3.")
    except ValueError:
        print("  Invalid input – please enter an integer.")

# ── geometry helpers ────────────────────────────────────

def bary_to_cart(lam, V):
    """Barycentric (3,) → Cartesian (2,) using triangle vertices V (3×2)."""
    return lam[0]*V[0] + lam[1]*V[1] + lam[2]*V[2]

# Equilateral triangle vertices (Cartesian)
V_cart = np.array([
    [0.5,  np.sqrt(3)/2],   # V1 – top
    [0.0,  0.0          ],   # V2 – lower left
    [1.0,  0.0          ],   # V3 – lower right
])

# ── build all triangles by recursive subdivision ────────────────────

# Barycentric vertices as exact Fractions
V1b = (Fraction(1), Fraction(0), Fraction(0))
V2b = (Fraction(0), Fraction(1), Fraction(0))
V3b = (Fraction(0), Fraction(0), Fraction(1))

def midpoint(A, B):
    return tuple((a + b) / 2 for a, b in zip(A, B))

def subdivide(tri):
    """Split triangle into 3 corner sub-triangles + 1 removed centre."""
    A, B, C = tri
    AB, BC, CA = midpoint(A, B), midpoint(B, C), midpoint(C, A)
    return [(A, AB, CA), (AB, B, BC), (CA, BC, C)], (AB, BC, CA)

def build_gasket(n):
    """Return (solid_list, removed_list) of triangles (barycentric) at level n."""
    solid   = [(V1b, V2b, V3b)]
    removed = []
    for _ in range(n):
        next_solid = []
        for tri in solid:
            corners, centre = subdivide(tri)
            next_solid.extend(corners)
            removed.append(centre)
        solid = next_solid
    return solid, removed   # 3^n solid,  (3^n - 1)/2 removed

solid_triangles, removed_triangles = build_gasket(n)

# ── draw the gasket ─────────────────────────────────────

fig, ax = plt.subplots(figsize=(7,6.5))
ax.set_aspect('equal')
ax.axis('off')
ax.set_xlim(-0.06, 1.06)
ax.set_ylim(-0.10, 0.99)

LIGHT_GRAY  = '#c8c8c8'
WHITE       = '#ffffff'
EDGE_COLOR  = '#444444'
POINT_COLOR = '#c0392b'

def draw_tri(tri_bary, fc, ec=EDGE_COLOR, lw=0.6, zorder=1):
    pts = np.array([
        bary_to_cart(np.array([float(v) for v in vert]), V_cart)
        for vert in tri_bary
    ])
    ax.add_patch(mpatches.Polygon(
        pts, closed=True, facecolor=fc, edgecolor=ec, linewidth=lw, zorder=zorder
    ))

# Outer boundary
draw_tri((V1b, V2b, V3b), fc=LIGHT_GRAY, lw=1.2, zorder=0)

# White removed regions (paint over the gray background)
for tri in removed_triangles:
    draw_tri(tri, fc=WHITE, ec=WHITE, lw=0, zorder=2)

# Solid small triangles with edges
for tri in solid_triangles:
    draw_tri(tri, fc=LIGHT_GRAY, ec=EDGE_COLOR, lw=max(0.3, 0.7 - 0.1*n), zorder=3)

# ── chaos game: all 3^(n+2) tuples of (n+2) moves ────────────────

num_moves = n + 2
D = 2 ** (n + 2)   # denominator

V_bary = {
    1: (Fraction(1), Fraction(0), Fraction(0)),
    2: (Fraction(0), Fraction(1), Fraction(0)),
    3: (Fraction(0), Fraction(0), Fraction(1)),
}

plotted_points = []   # Cartesian coords for plotting
bary_records  = []   # (moves, (a,b,c) integer numerators over D)
count = 0

for moves in iproduct([1, 2, 3], repeat=num_moves):
    # Starting contribution: initV / D
    coords = [Fraction(0), Fraction(0), Fraction(0)]
    coords[initV-1] = Fraction(1,D)

    # Move i (1-indexed): contribute V_{M[i]} / 2^(n+3-i)
    for i, m in enumerate(moves, start=1):
        factor = Fraction(1, 2 ** (n + 3 - i))
        vb = V_bary[m]
        for k in range(3):
            coords[k] += vb[k] * factor

    # Express as integer numerators over D
    int_nums = [int(c * D) for c in coords]
    assert sum(int_nums) == D, f"Coords don't sum to 1: {int_nums}"

    # Plot iff NO numerator is divisible by 4
    if all(x % 4 != 0 for x in int_nums):
        cart = bary_to_cart(np.array([float(c) for c in coords]), V_cart)
        plotted_points.append(cart)
        bary_records.append((moves, tuple(int_nums)))
        count += 1

# ── plot the points ────────────────────────────────────────

if plotted_points:
    pts = np.array(plotted_points)
    pt_size = 60 - 10 * n
    ax.scatter(pts[:, 0], pts[:, 1],
               s=pt_size, color=POINT_COLOR, zorder=5,
               linewidths=0, alpha=0.85)

# Vertex labels
for label, vc, ha in [('$V_1$', V_cart[0], 'center'),
                       ('$V_2$', V_cart[1], 'right'),
                       ('$V_3$', V_cart[2], 'left')]:
    dx = 0 if ha == 'center' else (-0.01 if ha == 'right' else 0.01)
    dy = 0.035 if ha == 'center' else -0.01
    ax.text(vc[0]+dx, vc[1]+dy, label, ha=ha, va='center', fontsize=14)

total_tuples = 3 ** num_moves
minMoves = 2+n
ax.set_title(
     f'Level-{n} Sierpinski Gasket with {minMoves} moves starting at $V_{initV}$\n'
     f'{count} interior points out of {total_tuples} possible points',
     fontsize=11, pad=10
 )

plt.tight_layout()
filename = f'sierpinski_level{n}-{initV}.png'
plt.savefig(filename, dpi=150, bbox_inches='tight')
plt.show()

print(f"\nLevel n            : {n}")
print(f"Number of moves    : {num_moves}  (= n + 2)")
print(f"Total tuples (3^{num_moves}) : {total_tuples:,}")
print(f"Solid triangles    : {len(solid_triangles):,}  (= 3^n)")
print(f"Points plotted     : {count:,}")
print(f"Saved to           : {filename}")