#!/usr/bin/env python3

"""This script sets up and initializes the design problem for a finite element analysis using dolfin adjoint based on SIMP method.

It defines the resolution, reads the design variables, and writes out the initial design to a file.
"""

from fenics import Constant
from fenics import FunctionSpace
from fenics import interpolate
from fenics import UnitCubeMesh
from fenics import XDMFFile
import numpy as np

from engibench.utils.cli import cast_argv

# Extract parameters
# NN: Resolution of the grid (arbitrary, affects performance)
# vol_f: Volume fraction for the control
NN, vol_f, output_path = cast_argv(int, float, str)

# Discretization step size (based on NN)
step = 1.0 / float(NN)

# Generate mesh grid values for both x, y and Z directions
x_values = np.linspace(0, 1, num=NN + 1)
y_values = np.linspace(0, 1, num=NN + 1)
z_values = np.linspace(0, 1, num=NN + 1)

# Set up the finite element function space
V = Constant(vol_f)  # Volume bound on the control
mesh = UnitCubeMesh(NN, NN, NN)  # Create a unit cube mesh with NN discretization
A = FunctionSpace(mesh, "CG", 1)  # Function space for the control variable

if __name__ == "__main__":
    # Initialize the design variable as a constant volume
    MM = V
    a = interpolate(MM, A)  # Initial guess for the design

    # Define the path to save the design files
    xdmf_file_path = output_path.removesuffix(".npy") + ".xdmf"

    # Write the mesh and initial design to an XDMF file
    with XDMFFile(xdmf_file_path) as outfile:
        outfile.write(mesh)  # Write mesh to the file
        outfile.write_checkpoint(a, "u", 0, append=True)  # Write design variable to the file

    # Initialize an array to store the results (x, y, z, value)
    results = np.zeros(((NN + 1) ** 3, 1))

    # Populate the results array with the mesh coordinates and the corresponding volume value
    ind = 0
    for _xs in x_values:
        for _ys in y_values:
            for _zs in z_values:
                results[ind, 0] = V  # Store the volume value
                ind += 1
    results = results.reshape(NN + 1, NN + 1, NN + 1)
    # Save the results array to a .npy file
    np.save(output_path, results)
