#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Apr 15 09:53:23 2021

@author: rmontant
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D   # needed for '3d" projection?

import heatsink    # import my file that defines the Heatsink class

# Produce a "scatter" plot of the coordinates of any object, viz. a heatsink:
def draw_heatsink(axis, heatsink):
    # define a reddish color
    def rgb(frac, transparency=0.2):
        return (np.sqrt(frac), frac, frac, transparency)
    #
    # Define each layer's color based on its depth (z-value).
    # This can be based on the initialized element values, as done in class,
    # or directly on the z-coordinate, as done here:
    color  = np.array( [rgb(0.75*v/heatsink.depth) for v in heatsink.coords[ : ,2]] )

    axis.scatter( \
        heatsink.coords[ : , 0],    # x / depth
        heatsink.coords[ : , 1],    # y / length
        heatsink.coords[ : , 2],    # z / depth
        color=color,                # z-dependent color
        marker=',',                 # use a small dot for each point
    )
#--------


# Build an array of 4 heatsinks, by means of a list comprehension:
# myhs = []
# for n in range(4):
#     myhs.append( heatsink.Heatsink(nfins=n+3) )
myhs = [ heatsink.Heatsink(nfins=n+3)  for n in range(4) ]

# Display some characteristics of the array of heatsinks:
for n in range(4):
    print('cube {:d}'.format(n), myhs[-1].cube.shape)
    print('     ', myhs[-1].coords.shape)


# Use the "draw_heatsink()" function to plot out each of the heatsinks,
# in a 2x2 grid of graphs within a figure:
fig = plt.figure()
axes = []
for idx in range(1, (4+1)):
    axis = fig.add_subplot(2, 2, idx, projection='3d')
    print(axis)
    draw_heatsink(axis, myhs[idx-1])
    axes.append(axis)

plt.show()  # go ahead and render the graph(s) in a display window somewhere