app.layout = html.Div([
    html.Div([
        dbc.InputGroup([
            dbc.Input(id="pdb-search-bar", type="text", placeholder="Enter PDB ID"),
            dbc.Button("Load", color="primary", id="load-button"),
        ]),
    ]),
    dashbio.Molecule3dViewer(
        id='pdb-viewer',
        modelData={'atoms': [], 'bonds': []},
        selectionType='residue',
        width='100%',
        style = {'display': 'flex'}
    )
])

@callback(
    Output(component_id='pdb-viewer', component_property='modelData'),
    Output(component_id='pdb-viewer', component_property='styles'),
    Output('pdb-viewer', 'selectedAtomIds'),
    Input(component_id='load-button', component_property='n_clicks'),
    State(component_id='pdb-search-bar', component_property='value'),
    prevent_initial_call=True
)
def update_pdb(n_clicks, pdb_id):
    data_update = Patch()
    styles_update = Patch()
    if not pdb_id:
        data_update.update({'atoms': [], 'bonds': []})
        styles_update.clear()
        selected_update = []
        return data_update, styles_update, selected_update
    else:
        try:
            parser = PdbParser(pdb_id)
            pdb_data = parser.mol3d_data()
        except:
            data_update.update({'atoms': [], 'bonds': []})
            styles_update.clear()
            selected_update = []
            return data_update, styles_update, selected_update
        color_element = 'chain'
        styles = create_mol3d_style(
            pdb_data['atoms'], visualization_type='cartoon', color_element=color_element
        )
        data_update.update(pdb_data)
        styles_update = styles
        selected_update = []
        return data_update, styles_update, selected_update