from pyintact import (
    MeshModel,
    IsotropicMaterialDescriptor,
    FixedBoundaryDescriptor,
    VectorForceDescriptor,
    LinearElasticScenarioDescriptor,
    StressSimulator,
    MaterialDomain,
    QueryResult,
    UnitSystem,
    FieldQuery,
    Field,
)
import os

# Setup paths
script_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(script_dir)

# Define aluminum material
aluminum = IsotropicMaterialDescriptor()
aluminum.density = 2700.0
aluminum.poisson_ratio = 0.33
aluminum.youngs_modulus = 69.0e9

# Define steel material
steel = IsotropicMaterialDescriptor()
steel.density = 27007800.0
steel.poisson_ratio = 0.3
steel.youngs_modulus = 2.1e11

# Load meshes
design_mesh = MeshModel("bracketassy_design_geometry.ply")
design_mesh.instance_id = "design_geometry"
design_mesh.refine()

nondesign_mesh = MeshModel("bracketassy_nondesign_geometry.ply")
nondesign_mesh.instance_id = "non_design_geometry"
nondesign_mesh.refine()

# Define boundary conditions
fixed_boundary = FixedBoundaryDescriptor()
fixed_boundary.boundary = MeshModel("bracketassy_fixed.ply")

vector_load = VectorForceDescriptor()
vector_load.boundary = MeshModel("bracketassy_load.ply")
vector_load.direction = [0.8728715609439696, -0.4364357804719848, -0.2182178902359924]
vector_load.magnitude = 10.0
vector_load.units = UnitSystem.MeterKilogramSecond

# Setup the simulation scenario descriptor
descriptor = LinearElasticScenarioDescriptor()
descriptor.materials = {"Aluminum 6061": aluminum, "Steel": steel}
descriptor.metadata.cell_size = 2.5
descriptor.metadata.tolerance = 1e-05
descriptor.boundary_conditions = [fixed_boundary, vector_load]

design_domain = MaterialDomain(design_mesh, aluminum)
nondesign_domain = MaterialDomain(nondesign_mesh, steel)
assembly = [design_domain, nondesign_domain]

# Initialize and run the simulator
simulator = StressSimulator(assembly, descriptor)
simulator.solve()

# Query results for stress distribution
results = QueryResult(assembly)
stress_query = FieldQuery(f=Field.VonMisesStress)

# Get the query result from the simulation
result = simulator.sample(stress_query, results)
results.writeVTK("multimaterial_assembly.vtu", UnitSystem.MeterKilogramSecond)