"""
PDB Writer
==========

"""

from __future__ import annotations

import pathlib
import typing

from stk._internal.molecule import Molecule
from stk._internal.periodic_info import PeriodicInfo
from stk._internal.utilities.utilities import OneOrMany


class PdbWriter:
    """
    A writer class for ``.pdb`` files.

    Examples:

        *Writing to a File with a Unit Cell*

        This writer can write to a file with the unit
        cell included for periodic molecules. Note that this always
        assumes P1 space group.

        .. testcode:: writing-to-a-file-with-a-unit-cell

            import stk

            bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
            bb2 = stk.BuildingBlock(
                smiles='BrCC(CBr)CBr',
                functional_groups=[stk.BromoFactory()],
            )
            topology_graph = stk.cof.PeriodicHoneycomb(
                building_blocks=(bb1, bb2),
                lattice_size=(3, 3, 1),
            )
            construction_result = topology_graph.construct()
            cof = (
                stk.ConstructedMolecule.init_from_construction_result(
                    construction_result=construction_result,
                )
            )
            writer = stk.PdbWriter()
            writer.write(
                molecule=cof,
                path='cof.pdb',
                periodic_info=construction_result.get_periodic_info(),
            )

        .. testcode:: writing-to-a-file-with-a-unit-cell
            :hide:

            import os

            assert os.path.exists('cof.pdb')

        .. testcleanup:: writing-to-a-file-with-a-unit-cell

            os.remove('cof.pdb')

    """

    def _write_content(
        self,
        molecule: Molecule,
        atom_ids: typing.Optional[OneOrMany[int]],
        periodic_info: typing.Optional[PeriodicInfo] = None,
    ) -> list[str]:
        if atom_ids is None:
            atom_ids = range(molecule.get_num_atoms())
        elif isinstance(atom_ids, int):
            atom_ids = (atom_ids,)

        content = []
        if periodic_info is not None:
            # Input unit cell information.
            a = periodic_info.get_a()
            b = periodic_info.get_b()
            c = periodic_info.get_c()
            alpha = periodic_info.get_alpha()
            beta = periodic_info.get_beta()
            gamma = periodic_info.get_gamma()
            content.append(
                f"CRYST1 {a:>8.3f} {b:>8.3f} {c:>8.3f}"
                f" {alpha:>6.2f} {beta:>6.2f} {gamma:>6.2f} "
                f"P 1\n"
            )

        atom_counts: dict[str, int] = {}
        hetatm = "HETATM"
        alt_loc = ""
        res_name = "UNL"
        chain_id = ""
        res_seq = "1"
        i_code = ""
        occupancy = "1.00"
        temp_factor = "0.00"

        coords = molecule.get_position_matrix()
        # This set will be used by bonds.
        atoms = set()
        for atom_id in atom_ids:
            atoms.add(atom_id)
            (atom,) = molecule.get_atoms(atom_ids=atom_id)
            serial = atom_id + 1
            element = atom.__class__.__name__
            charge = atom.get_charge()
            atom_counts[element] = atom_counts.get(element, 0) + 1
            name = f"{element}{atom_counts[element]}"
            # Make sure the coords are no more than 8 columns wide
            # each.
            x, y, z = (i for i in coords[atom_id])

            content.append(
                f"{hetatm:<6}{serial:>5} {name:<4}"
                f"{alt_loc:<1}{res_name:<3} {chain_id:<1}"
                f"{res_seq:>4}{i_code:<1}   "
                f" {x:>7.3f} {y:>7.3f} {z:>7.3f}"
                f"{occupancy:>6}{temp_factor:>6}          "
                f"{element:>2}{charge:>2}\n"
            )

        conect = "CONECT"
        for bond in molecule.get_bonds():
            a1 = bond.get_atom1().get_id()
            a2 = bond.get_atom2().get_id()
            if a1 in atoms and a2 in atoms:
                content.append(
                    f"{conect:<6}{a1 + 1:>5}{a2 + 1:>5}               \n"
                )

        content.append("END\n")

        return content

    def to_string(
        self,
        molecule: Molecule,
        atom_ids: typing.Optional[OneOrMany[int]] = None,
        periodic_info: typing.Optional[PeriodicInfo] = None,
    ) -> str:
        """
        Get a ``.pdb`` file format string of `molecule`.

        Parameters:

            molecule:
                Molecule to write to ``.pdb`` format.

            atom_ids:
                The atom ids of atoms to write. Can be a single
                :class:`int`, if a single atom is to be used, or
                ``None``, if all atoms are to be used.

            periodic_info:
                Information about the periodic cell.

        Returns:

            A string holding the content of a ``.pdf`` file.

        """

        content = self._write_content(
            molecule=molecule,
            atom_ids=atom_ids,
            periodic_info=periodic_info,
        )

        return "".join(content)

    def write(
        self,
        molecule: Molecule,
        path: pathlib.Path | str,
        atom_ids: typing.Optional[OneOrMany[int]] = None,
        periodic_info: typing.Optional[PeriodicInfo] = None,
    ) -> None:
        """
        Write `molecule` to ``.pdb`` file format.

        Parameters:

            molecule:
                Molecule to write to ``.pdb`` format.

            path:
                The full path to the file being written.

            atom_ids:
                The atom ids of atoms to write. Can be a single
                :class:`int`, if a single atom is to be used, or
                ``None``, if all atoms are to be used.

            periodic_info:
                Information about the periodic cell.

        """

        content = self._write_content(
            molecule=molecule,
            atom_ids=atom_ids,
            periodic_info=periodic_info,
        )

        with open(path, "w") as f:
            f.write("".join(content))
