import json
import sqlite3

class Molecule:
    def __init__(self, name, multiplicity, geometry, properties):
        self.name = name
        self.multiplicity = multiplicity
        self.geometry = geometry
        self.properties = properties

    def to_dict(self):
        return {
            "name": self.name,
            "multiplicity": self.multiplicity,
            "geometry": self.geometry,
            "properties": self.properties,
        }

    @staticmethod
    def from_dict(data):
        return Molecule(
            name=data["name"],
            multiplicity=data["multiplicity"],
            geometry=data["geometry"],
            properties=data["properties"]
        )

def save_molecules_to_json(molecules, filename):
    with open(filename, 'w') as f:
        json_data = [molecule.to_dict() for molecule in molecules]
        json.dump(json_data, f, indent=4)

def load_molecules_from_json(filename):
    with open(filename, 'r') as f:
        json_data = json.load(f)
        return [Molecule.from_dict(data) for data in json_data]


def create_database(db_name):
    conn = sqlite3.connect(db_name)
    cursor = conn.cursor()
    cursor.execute('''CREATE TABLE IF NOT EXISTS molecules
                      (name TEXT, multiplicity INTEGER, geometry TEXT, properties TEXT)''')
    conn.commit()
    conn.close()

def add_molecule_to_db(db_name, molecule):
    conn = sqlite3.connect(db_name)
    cursor = conn.cursor()
    geometry_str = json.dumps(molecule.geometry)
    energies_str = json.dumps(molecule.properties)
    cursor.execute("INSERT INTO molecules VALUES (?, ?, ?, ?)",
                   (molecule.name, molecule.multiplicity, geometry_str, energies_str))
    conn.commit()
    conn.close()

def get_molecules_from_db(db_name):
    conn = sqlite3.connect(db_name)
    cursor = conn.cursor()
    cursor.execute("SELECT name, multiplicity, geometry, properties FROM molecules")
    rows = cursor.fetchall()
    molecules = []
    for row in rows:
        name, multiplicity, geometry_str, energies_str = row
        geometry = json.loads(geometry_str)
        properties = json.loads(energies_str)
        molecules.append(Molecule(name, multiplicity, geometry, properties))
    conn.close()
    return molecules


def generate_xyz(elements, filename="output.xyz", verbose=False):
    """
    Generate an XYZ file from a list of elements.

    Parameters:
    elements (list): A list of dictionaries, where each dictionary represents
                     an atom with its element and x, y, z coordinates.
    filename (str): The name of the output XYZ file. Default is 'output.xyz'.
    """

    # Get the number of atoms
    num_atoms = len(elements)

    # Open the file in write mode
    with open(filename, 'w') as f:
        # Write the number of atoms
        f.write(f"{num_atoms}\n")

        # Write a comment line (can be left blank or customized)
        f.write("XYZ file generated by generate_xyz function\n")

        # Write the element and coordinates
        for atom in elements:
            element = atom['element']
            x = atom['x']
            y = atom['y']
            z = atom['z']
            f.write(f"{element} {x:.6f} {y:.6f} {z:.6f}\n")

    if(verbose):
        print(f"XYZ file '{filename}' generated successfully!")