#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#------------------------------------------------
#  Purpose: Volume generation of hexagonal arrays
#   Author: Jeremie Vassseur
#    Email: jeremie.vasseur@min.uni-muenchen.de
#    Email: vasseur.jeremie@gmail.com
#
# Copyright (C) 2019 Jeremie Vasseur
#------------------------------------------------

import os
import numpy as np
from cv2 import imwrite

# calculates the normalized radius of a 'n'-gon (i.e. the distance between
# its center and its border) for any angle 'theta' ranging from 0 to 2*np.pi
calc_r_ngon = lambda theta, n: np.cos(np.pi/n)/np.cos(theta%(2*np.pi/n)-np.pi/n)

# calculates the dimensions of a 2D hexagonal array containing 'nx' and
# 'ny' hexagons with side length 'b' in the x- and y-direction respectively
calc_dim = lambda b, nx, ny: np.array((1.5*nx*b+1, np.round(ny*b*np.sqrt(3))+1), int)

# constructs the 2D array with fracture width 'w'
def build_data_2d(b, w, nx, ny):
    dim = calc_dim(b, nx, ny)
    data = np.ones(dim[::-1], int)
    x0, y0 = 0, 0
    for ic in range(nx+1):
        for jc in range(ny+1):
            if ic%2==0: k = 0
            else: k = 1/2
            x, y = int(x0+1.5*ic*b), int(np.round(y0+b*np.sqrt(3)*(jc+k)))
            for i in range(2*b+1):
                for j in range(2*b+1):
                    dx, dy = i-b, j-b
                    xc, yc = x+dx, y+dy
                    if dx==0: theta = np.pi/2
                    else: theta = np.arctan2(dy, dx)
                    r = (b-w)*calc_r_ngon(theta, 6)
                    if (xc>=0 and xc<dim[0] and yc>=0 and yc<dim[1]
                        and dx**2+dy**2<=r**2 and data[yc,xc]!=0):
                        data[yc,xc] = 0
    return data

# constructs the 3D arrays for a given set of parameter and saves it as a stack
# of tif images which can then be fed into LBflow for fluid flow characterization
b, nx, ny = [50, 100], [8, 4], [6, 3] # 'b' is given in pixels
for i in range(len(b)):
    w = np.arange(1, int(b[i]/5)+1, 1) # 'w' is also given in pixels
    for j in range(len(w)):
        data = build_data_2d(b[i], w[j], nx[i], ny[i])
        dim = np.append(data.shape, 50)
        name = "b_{:d}_w_{:d}_nx_{:d}_ny_{:d}_dim_{}".format(
            b[i], w[j], nx[i], ny[i], "{:d}_{:d}_{:d}".format(*dim))
        if not os.path.exists(name):
            os.makedirs(name)
            for k in range(dim[2]):
                fname = "hex_array_{}_{:04d}.tif".format(name, k+1)
                imwrite(os.path.join(name, fname), 255*data)