from plotlymol3d import draw_3D_rep

# Draw a molecule from SMILES
fig = draw_3D_rep(smiles="CCNCOCSC", mode="ball+stick", ambient=0.1)
fig.show()

# Draw from XYZ file
fig = draw_3D_rep(xyzfile="path/to/file.xyz", mode="ball+stick", ambient=0.1)
fig.show()

from plotlymol3d import draw_3D_rep

fig = draw_3D_rep(
    cubefile="path/to/file.cube",
    molfile="path/to/file.mol",
    mode="ball+stick",
    ambient=0.1,
    cubedraw="orbitals",
    orbital_opacity=0.25,
    orbital_colors=["darkorange", "darkblue"],
)
fig.show()

from plotlymol3d import draw_3D_rep

fig = draw_3D_rep(
    smiles="O",  # Water molecule
    vibration_file="water_freq.log",
    vibration_mode=1,  # First vibrational mode
    vibration_display="arrows",
    vibration_amplitude=1.5,
)
fig.show()

from plotlymol3d import parse_vibrations, create_vibration_animation
from rdkit.Chem import MolFromSmiles, AddHs
from rdkit.Chem.AllChem import EmbedMolecule

# Parse vibration data
vib_data = parse_vibrations("water_freq.log")

# Create molecule
mol = MolFromSmiles("O")
mol = AddHs(mol)
EmbedMolecule(mol)

# Generate animation
fig = create_vibration_animation(
    vib_data=vib_data,
    mode_number=1,
    mol=mol,
    amplitude=0.5,
    n_frames=20,  # Smoother with more frames
    mode="ball+stick"
)
fig.show()

from plotlymol3d import draw_3D_rep, parse_vibrations, add_vibrations_to_figure

# Create molecular figure
fig = draw_3D_rep(smiles="O", mode="ball+stick")

# Parse vibrations and add heatmap
vib_data = parse_vibrations("water_freq.log")
fig = add_vibrations_to_figure(
    fig=fig,
    vib_data=vib_data,
    mode_number=1,
    display_type="heatmap",
    heatmap_colorscale="Reds"
)
fig.show()

from plotlymol3d import (
    parse_gaussian_vibrations,  # Gaussian .log files
    parse_orca_vibrations,       # ORCA .out files
    parse_molden_vibrations,     # Molden .molden files
    parse_vibrations,            # Auto-detect format
)

# Auto-detect format from file extension
vib_data = parse_vibrations("calculation.log")

# Access mode data
for mode in vib_data.modes:
    print(f"Mode {mode.mode_number}: {mode.frequency:.1f} cm⁻¹")
    if mode.ir_intensity:
        print(f"  IR Intensity: {mode.ir_intensity:.1f} km/mol")