from sgio import read, readOutputState, write, addPointDictDataToMesh
from sgio.iofunc.gmsh import write_element_node_data

fn_in = 'wing3dSG_adjusted.sc'
ff_in = 'sc'
fv_in = '2.1'
fn_out = 'wing3dSG_adjusted.msh'
ff_out = 'gmsh'
fv_out = '4.1'
sgdim = 3
model_type = 'bm2'
lout_exts = ['u', 'sn', 'snm']

# Read the cross-section model
sg = read(fn_in, ff_in, sgdim=sgdim, model_type=model_type)

# Read the output state
state_cases = readOutputState(
    fn_in, ff_in, 'd',
    extension=lout_exts, sg=sg, tool_version=fv_in,
    num_cases=1, output_format=1
    )

print(state_cases[0].states.keys())

for j, state_case in enumerate(state_cases):
    addPointDictDataToMesh(['u1', 'u2', 'u3'], state_case.getState('u').data, sg.mesh)

# Write the mesh to a file
write(sg, fn_out, ff_out, format_version=fv_out)

# Write the output state to the mesh file
for j, state_case in enumerate(state_cases):
    for _name, _state in state_case.states.items():
        if _state.location == 'element_node':
            with open(fn_out, 'a') as file:
                write_element_node_data(file, _state.label[0], _state.data)