__name__ = "fastpdb"
__author__ = "Patrick Kunzmann"
__all__ = ["PDBFile"]
__version__ = "1.0.0"

import numpy as np
import biotite
import biotite.structure as struc
import biotite.structure.io.pdb as pdb
from .fastpdb import PDBFile as RustPDBFile


class PDBFile(biotite.TextFile):
    r"""
    This class represents a PDB file.
    
    This class only provides support for reading/writing the pure atom
    information (``ATOM``, ``HETATM``, ``MODEL`` and ``ENDMDL``
    records).
    ``TER`` records cannot be written.
    
    See also
    --------
    PDBxFile
    
    Examples
    --------
    Load a ``\\*.pdb`` file, modify the structure and save the new
    structure into a new file:
    
    >>> import os.path
    >>> file = PDBFile.read(os.path.join(path_to_structures, "1l2y.pdb"))
    >>> array_stack = file.get_structure()
    >>> array_stack_mod = rotate(array_stack, [1,2,3])
    >>> file = PDBFile()
    >>> file.set_structure(array_stack_mod)
    >>> file.write(os.path.join(path_to_directory, "1l2y_mod.pdb"))
    """

    def __init__(self):
        super().__init__()
        self._pdb_file = RustPDBFile([])
    
    @classmethod
    def read(cls, file):
        file = super().read(file)
        file._pdb_file = RustPDBFile(file.lines)
        return file
    
    def get_model_count(self):
        """
        Get the number of models contained in the PDB file.

        Returns
        -------
        model_count : int
            The number of models.
        """
        return self._pdb_file.get_model_count()

    def get_coord(self, model=None):
        """
        Get only the coordinates of the PDB file.
        
        Parameters
        ----------
        model : int, optional
            If this parameter is given, the function will return a
            2D coordinate array from the atoms corresponding to the
            given model number (starting at 1).
            Negative values are used to index models starting from the
            last model insted of the first model.
            If this parameter is omitted, an 2D coordinate array
            containing all models will be returned, even if
            the structure contains only one model.
        
        Returns
        -------
        coord : ndarray, shape=(m,n,3) or shape=(n,2), dtype=float
            The coordinates read from the ``ATOM`` and ``HETATM``
            records of the file.
        
        Notes
        -----
        Note that :func:`get_coord()` may output more coordinates than
        the atom array (stack) from the corresponding
        :func:`get_structure()` call has.
        The reason for this is, that :func:`get_structure()` filters
        *altloc* IDs, while `get_coord()` does not.
        
        Examples
        --------
        Read an :class:`AtomArrayStack` from multiple PDB files, where
        each PDB file contains the same atoms but different positions.
        This is an efficient approach when a trajectory is spread into
        multiple PDB files, as done e.g. by the *Rosetta* modeling
        software. 

        For the purpose of this example, the PDB files are created from
        an existing :class:`AtomArrayStack`.
        
        >>> import os.path
        >>> from tempfile import gettempdir
        >>> file_names = []
        >>> for i in range(atom_array_stack.stack_depth()):
        ...     pdb_file = PDBFile()
        ...     pdb_file.set_structure(atom_array_stack[i])
        ...     file_name = os.path.join(gettempdir(), f"model_{i+1}.pdb")
        ...     pdb_file.write(file_name)
        ...     file_names.append(file_name)
        >>> print(file_names)
        ['...model_1.pdb', '...model_2.pdb', ..., '...model_38.pdb']

        Now the PDB files are used to create an :class:`AtomArrayStack`,
        where each model represents a different model.

        Construct a new :class:`AtomArrayStack` with annotations taken
        from one of the created files used as template and coordinates
        from all of the PDB files.

        >>> template_file = PDBFile.read(file_names[0])
        >>> template = template_file.get_structure()
        >>> coord = []
        >>> for i, file_name in enumerate(file_names):
        ...     pdb_file = PDBFile.read(file_name)
        ...     coord.append(pdb_file.get_coord(model=1))
        >>> new_stack = from_template(template, np.array(coord))

        The newly created :class:`AtomArrayStack` should now be equal to
        the :class:`AtomArrayStack` the PDB files were created from.

        >>> print(np.allclose(new_stack.coord, atom_array_stack.coord))
        True
        """
        if model is None:
            coord = self._pdb_file.parse_coord_multi_model()
        else:
            coord = self._pdb_file.parse_coord_single_model(model)
        return coord
    
    def get_structure(self, model=None, altloc="first", extra_fields=None, include_bonds=False):
        """
        Get an :class:`AtomArray` or :class:`AtomArrayStack` from the PDB file.
        
        Parameters
        ----------
        model : int, optional
            If this parameter is given, the function will return an
            :class:`AtomArray` from the atoms corresponding to the given
            model number (starting at 1).
            Negative values are used to index models starting from the
            last model insted of the first model.
            If this parameter is omitted, an :class:`AtomArrayStack`
            containing all models will be returned, even if the
            structure contains only one model.
        altloc : {'first', 'occupancy', 'all'}
            This parameter defines how *altloc* IDs are handled:
                - ``'first'`` - Use atoms that have the first
                  *altloc* ID appearing in a residue.
                - ``'occupancy'`` - Use atoms that have the *altloc* ID
                  with the highest occupancy for a residue.
                - ``'all'`` - Use all atoms.
                  Note that this leads to duplicate atoms.
                  When this option is chosen, the ``altloc_id``
                  annotation array is added to the returned structure.
        extra_fields : list of str, optional
            The strings in the list are optional annotation categories
            that should be stored in the output array or stack.
            These are valid values:
            ``'atom_id'``, ``'b_factor'``, ``'occupancy'`` and
            ``'charge'``.
        include_bonds : bool, optional
            If set to true, a :class:`BondList` will be created for the
            resulting :class:`AtomArray` containing the bond information
            from the file.
            All bonds have :attr:`BondType.ANY`, since the PDB format
            does not support bond orders.
        
        Returns
        -------
        array : AtomArray or AtomArrayStack
            The return type depends on the `model` parameter.
        """
        if extra_fields is not None:
            include_atom_id   = "atom_id"   in extra_fields
            include_b_factor  = "b_factor"  in extra_fields
            include_occupancy = "occupancy" in extra_fields
            include_charge    = "charge"    in extra_fields
        else:
            include_atom_id   = False
            include_b_factor  = False
            include_occupancy = False
            include_charge    = False
        if include_bonds:
            # Required for mapping the bonded atom IDs to atom indices
            include_atom_id = True
        if altloc == "occupancy":
            include_occupancy = True

        
        if model is None:
            coord = self._pdb_file.parse_coord_multi_model()
            annotations = self._pdb_file.parse_annotations(
                1,
                include_atom_id, include_b_factor,
                include_occupancy, include_charge
            )
        else:
            coord = self._pdb_file.parse_coord_single_model(model)
            annotations = self._pdb_file.parse_annotations(
                model,
                include_atom_id, include_b_factor,
                include_occupancy, include_charge
            )
        (
            chain_id, res_id, ins_code, res_name,
            hetero, atom_name, element, altloc_id,
            atom_id, b_factor, occupancy, charge
        ) = annotations
        # Interpret uint32 arrays as unicode arrays
        chain_id  = np.frombuffer(chain_id,  dtype="U4")
        ins_code  = np.frombuffer(ins_code,  dtype="U1")
        res_name  = np.frombuffer(res_name,  dtype="U3")
        atom_name = np.frombuffer(atom_name, dtype="U6")
        element   = np.frombuffer(element,   dtype="U2")
        altloc_id = np.frombuffer(altloc_id, dtype="U1")
        
        if coord.ndim == 3:
            atoms = struc.AtomArrayStack(coord.shape[0], coord.shape[1])
            atoms.coord = coord
        else:
            atoms = struc.AtomArray(coord.shape[0])
            atoms.coord = coord

        atoms.chain_id  = chain_id
        atoms.res_id    = res_id
        atoms.ins_code  = ins_code
        atoms.res_name  = res_name
        atoms.hetero    = hetero
        atoms.atom_name = atom_name
        atoms.element   = element
        
        for field in (extra_fields if extra_fields is not None else []):
            if field == "atom_id":
                # Copy is necessary to avoid double masking in 
                # later altloc ID filtering
                atoms.set_annotation("atom_id", atom_id.copy())
            elif field == "charge":
                atoms.set_annotation("charge", charge)
            elif field == "occupancy":
                atoms.set_annotation("occupancy", occupancy)
            elif field == "b_factor":
                atoms.set_annotation("b_factor", b_factor)
            else:
                raise ValueError(f"Unknown extra field: {field}")

        
        box = self._pdb_file.parse_box()
        if box is None:
            atoms.box = None
        else:
            len_a, len_b, len_c, alpha, beta, gamma = box
            box = struc.vectors_from_unitcell(
                len_a, len_b, len_c,
                np.deg2rad(alpha), np.deg2rad(beta), np.deg2rad(gamma)
            )
        if isinstance(atoms, struc.AtomArray):
            atoms.box = box
        else:
            atoms.box = np.repeat(
                box[np.newaxis, ...], atoms.stack_depth(), axis=0
            )
        

        # Filter altloc IDs
        if altloc == "occupancy":
            filter = struc.filter_highest_occupancy_altloc(
                atoms, altloc_id, occupancy
            )
            atoms = atoms[..., filter]
            atom_id = atom_id[filter] if atom_id is not None else None
        elif altloc == "first":
            filter = struc.filter_first_altloc(atoms, altloc_id)
            atoms = atoms[..., filter]
            atom_id = atom_id[filter] if atom_id is not None else None
        elif altloc == "all":
            atoms.set_annotation("altloc_id", altloc_id)
        else:
            raise ValueError(f"'{altloc}' is not a valid 'altloc' option")
        

        if include_bonds:
            bond_list = struc.BondList(
                atoms.array_length(), self._pdb_file.parse_bonds(atom_id)
            )
            bond_list = bond_list.merge(struc.connect_via_residue_names(
                atoms,
                # The information for non-hetero residues and water
                # are not part of CONECT records
                (~atoms.hetero) | struc.filter_solvent(atoms)
            ))
            # Remove bond order from inter residue bonds for consistency
            bond_list.remove_bond_order()
            atoms.bonds = bond_list


        return atoms
    

    def set_structure(self, atoms):
        """
        Set the :class:`AtomArray` or :class:`AtomArrayStack` for the
        file.
        
        This makes also use of the optional annotation arrays
        ``'atom_id'``, ``'b_factor'``, ``'occupancy'`` and ``'charge'``.
        If the atom array (stack) contains the annotation ``'atom_id'``,
        these values will be used for atom numbering instead of
        continuous numbering.
        
        Parameters
        ----------
        array : AtomArray or AtomArrayStack
            The array or stack to be saved into this file. If a stack
            is given, each array in the stack is saved as separate
            model.
        
        Notes
        -----
        If `array` has an associated :class:`BondList`, ``CONECT``
        records are also written for all non-water hetero residues
        and all inter-residue connections.
        """
        # Reset lines of text
        self._pdb_file = RustPDBFile([])


        # Write 'CRYST1' record
        if atoms.box is not None:
            box = atoms.box
            if box.ndim == 3:
                box = box[0]
            len_a, len_b, len_c, alpha, beta, gamma \
                = struc.unitcell_from_vectors(box)
            self._pdb_file.write_box(
                len_a, len_b, len_c,
                np.rad2deg(alpha), np.rad2deg(beta), np.rad2deg(gamma)
            )
        
        
        # Write 'ATOM' and 'MODEL' records
        # Convert Unicode arrays into uint32 arrays for usage in Rust
        chain_id  = np.frombuffer(atoms.chain_id,  dtype=np.uint32).reshape(-1, 4)
        ins_code  = np.frombuffer(atoms.ins_code,  dtype=np.uint32).reshape(-1, 1)
        res_name  = np.frombuffer(atoms.res_name,  dtype=np.uint32).reshape(-1, 3)
        atom_name = np.frombuffer(atoms.atom_name, dtype=np.uint32).reshape(-1, 6)
        element   = np.frombuffer(atoms.element,   dtype=np.uint32).reshape(-1, 2)
        
        categories = atoms.get_annotation_categories()
        atom_id   = atoms.atom_id   if "atom_id"   in categories else None
        b_factor  = atoms.b_factor  if "b_factor"  in categories else None
        occupancy = atoms.occupancy if "occupancy" in categories else None
        charge    = atoms.charge    if "charge"    in categories else None

        # Convert to correct dtype for Rust function call, if necessary
        coord = atoms.coord.astype(np.float32, copy=False)
        res_id = atoms.res_id.astype(np.int64, copy=False)
        hetero = atoms.hetero.astype(bool, copy=False)
        if atom_id is not None:
            atom_id = atom_id.astype(np.int64, copy=False)
        if b_factor is not None:
            b_factor = b_factor.astype(np.float64, copy=False)
        if occupancy is not None:
            occupancy = occupancy.astype(np.float64, copy=False)
        if charge is not None:
            charge = charge.astype(np.int64, copy=False)
        
        if isinstance(atoms, struc.AtomArray):
            self._pdb_file.write_single_model(
                coord, chain_id, res_id, ins_code,
                res_name, hetero, atom_name, element,
                atom_id, b_factor, occupancy, charge
            )
        elif isinstance(atoms, struc.AtomArrayStack):
            self._pdb_file.write_multi_model(
                coord, chain_id, res_id, ins_code,
                res_name, hetero, atom_name, element,
                atom_id, b_factor, occupancy, charge
            )
        else:
            raise TypeError(
                f"Expected AtomArray or AtomArrayStack, "
                f"but got {type(atoms).__name__}"
            )

       
        # Write 'CONECT' records
        if atoms.bonds is not None:
            # Only non-water hetero records and connections between
            # residues are added to the records
            hetero_indices = np.where(atoms.hetero & ~struc.filter_solvent(atoms))[0]
            bond_array = atoms.bonds.as_array()
            bond_array = bond_array[
                np.isin(bond_array[:,0], hetero_indices) |
                np.isin(bond_array[:,1], hetero_indices) |
                (atoms.res_id  [bond_array[:,0]] != atoms.res_id  [bond_array[:,1]]) |
                (atoms.chain_id[bond_array[:,0]] != atoms.chain_id[bond_array[:,1]])
            ]
            # Bond type is unused since PDB does not support bond orders
            bonds, _ = struc.BondList(
                atoms.array_length(), bond_array
            ).get_all_bonds()
            atom_id = np.arange(1, atoms.array_length()+1, dtype=np.int64) \
                      if atom_id is None else atom_id
            self._pdb_file.write_bonds(
                bonds.astype(np.int32, copy=False), atom_id
            )

            
        self.lines = self._pdb_file.lines