import os

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

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

# Load L-bracket geometry
geometry = MeshModel("L-bracket.stl")
geometry.refine(0.02)
geometry.instance_id = "L_bracket"

# Define boundary conditions
fixed_bc1 = FixedBoundaryDescriptor()
fixed_bc1.boundary = MeshModel("L-bracket-bc1.stl")

fixed_bc2 = FixedBoundaryDescriptor()
fixed_bc2.boundary = MeshModel("L-bracket-bc2.stl")

load = VectorForceDescriptor()
load.boundary = MeshModel("L-bracket-load.stl")
load.direction = [0.0, 0.0, -1.0]
load.magnitude = 1000.0
load.units = UnitSystem.MeterKilogramSecond

# Define aluminum material
material = IsotropicMaterialDescriptor()
material.density = 2700.0  # kg/m^3
material.poisson_ratio = 0.33
material.youngs_modulus = 69.0e9  # Pa

# Setup the simulation scenario
descriptor = LinearElasticScenarioDescriptor()
descriptor.materials = {"Aluminum": material}
descriptor.boundary_conditions = [fixed_bc1, fixed_bc2, load]
descriptor.metadata.resolution = 20000
descriptor.metadata.units = UnitSystem.MeterKilogramSecond
descriptor.metadata.solver_override = Solver.AMGCL_amg_rigid_body

# Enable Adaptive Mesh Refinement (AMR) - Beta
# AMR is only compatible with linear elasticity.
descriptor.metadata.integrator_override = Integrator.AdaptiveLinearElasticity
descriptor.metadata.amr.refine_energy_threshold = 0.2
descriptor.metadata.amr.max_iterations = 8
descriptor.metadata.amr.max_refinement_depth = 4
descriptor.metadata.amr.max_resolution = 5000000

domain = MaterialDomain(geometry, "Aluminum", descriptor)
assembly = [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)
result = simulator.sample(stress_query, results)
results.writeVTK("amr.vtu", UnitSystem.MeterKilogramSecond)