From cf1786f6c1c351060c6206d37936482317b71bbb Mon Sep 17 00:00:00 2001
From: CalCraven <nicholas.c.craven@vanderbilt.edu>
Date: Thu, 9 Jan 2025 16:25:57 -0600
Subject: [PATCH] Add more structure for mol2 writer

---
 gmso/formats/mol2.py | 95 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 93 insertions(+), 2 deletions(-)

diff --git a/gmso/formats/mol2.py b/gmso/formats/mol2.py
index 89caceb7..aa050b42 100644
--- a/gmso/formats/mol2.py
+++ b/gmso/formats/mol2.py
@@ -3,13 +3,14 @@
 import itertools
 import warnings
 from pathlib import Path
+import os
 
 import unyt as u
 
 from gmso import Atom, Bond, Box, Topology
 from gmso.abc.abstract_site import Molecule, Residue
 from gmso.core.element import element_by_name, element_by_symbol
-from gmso.formats.formats_registry import loads_as
+from gmso.formats.formats_registry import loads_as, saves_as
 
 
 @loads_as(".mol2")
@@ -88,7 +89,69 @@ def read_mol2(filename, site_type="atom", verbose=False):
     return topology
 
 @saves_as(".mol2")
-def write_mol2()
+def write_mol2(top, filename):
+    """Write a gmso.Topology to a TRIPOS mol2 file. 
+
+    Creates a mol2 file from a Topology structure. This will generate the following
+    sections in the mol2 file.
+    - @<TRIPOS>ATOM
+    - @<TRIPOS>BOND
+    ??- @<TRIPOS>CRYSIN??
+    - @<TRIPOS>FF_PBC
+    - @<TRIPOS>MOLECULE
+
+    Parameters
+    ----------
+    filename : string
+        path to the file where the mol2 file will be saved
+    overwrite : bool, optional, default=False
+        If True, overwrite the file if it already exists. Raise error if False and 
+        overwriting is attempted.
+    site_type : string ('atom' or 'lj'), default='atom'
+        tells the reader to consider the elements saved in the mol2 file, and
+        if the type is 'lj', to not try to identify the element of the site,
+        instead saving the site name.
+    verbose : bool, optional, default=False
+        If True, raise warnings for any assumptions made during the parsing.
+
+    Returns
+    -------
+    None
+    """
+    # Massaging/type checking
+
+    # Loop through topology and write each section
+    with open(filename, "w") as f:
+        
+
+        #ATOM top.sites
+        #@<TRIPOS>ATOM
+        # 1 C          -0.7600    1.1691   -0.0005 C.ar    1  BENZENE       0.000
+        f.writelines("@<TRIPOS>ATOM\n")
+        for site in top.sites:
+            _write_site_info(site, f)
+
+
+    #BOND top.bonds
+    # bond_id origin_atom_id target_atom_id bond_type
+        for bond in top.bonds:
+            _write_bond_info(bond, f)
+
+    #FF_PBC
+    #format_version_number pbc_type pbc_x_coord_min
+    # pbc_y_coord_min pbc_z_coord_min pbc_x_coord_max
+    # pbc_y_coord_max pbc_z_coord_max solvent_type
+    # num_solvent_shells reorient_molecule_flag
+    # status_flag apply_
+        if top.box:
+            _write_box_info(top.box, f)
+
+    #MOLECULE 
+    # mol_name
+    # num_atoms [num_bonds [num_subst [num_feat[num_sets]]]]
+    # mol_type
+    # charge_type
+    _write_molecule_info(top, f)
 
 def _parse_lj(top, section, verbose):
     """Parse atom of lj style from mol2 file."""
@@ -197,3 +260,31 @@ def _parse_box(top, section, verbose):
 def _parse_molecule(top, section, verbose):
     """Parse molecule information from the mol2 file."""
     top.label = str(section[0].strip())
+
+def _write_site_info(site, f, index=1):
+    """Write site information to ATOM section."""
+    # TODO: Create rules to make sure nothing is too long, so that it cuts off.
+    lineList = [
+        str(index), 
+        site.element.symbol, 
+        *map(str,site.position.value), 
+        site.atom_type.name if site.atom_type else site.name, 
+        str(site.molecule.number), 
+        site.molecule.name, 
+        str(site.charge) if site.charge else "0.00"
+    ]
+    formattedStr = "\t".join(lineList)+"\n"
+    f.writelines(formattedStr)
+    index+=1
+    #ATOM top.sites
+    #@<TRIPOS>ATOM
+    # 1 C          -0.7600    1.1691   -0.0005 C.ar    1  BENZENE       0.000
+
+def _write_bond_info(bond, f):
+    pass
+
+def _write_molecule_info(top, f):
+    pass
+
+def _write_box_info(box, f):
+    pass