"""Utilities for handling PDBQT files."""

from typing import Dict, List, Optional, Set, Tuple
from deepchem.utils.typing import RDKitMol


def pdbqt_to_pdb(filename: Optional[str] = None,
                 pdbqt_data: Optional[List[str]] = None) -> str:
    """Extracts the PDB part of a pdbqt file as a string.

    Either `filename` or `pdbqt_data` must be provided. This function
    strips PDBQT charge information from the provided input.

    Parameters
    ----------
    filename: str, optional  (default None)
        Filename of PDBQT file
    pdbqt_data: List[str], optional (default None)
        Raw list of lines containing data from PDBQT file.

    Returns
    -------
    pdb_block: str
        String containing the PDB portion of pdbqt file.
    """
    if filename is not None and pdbqt_data is not None:
        raise ValueError("Only one of filename or pdbqt_data can be provided")
    elif filename is None and pdbqt_data is None:
        raise ValueError("Either filename or pdbqt_data must be provided")
    elif filename is not None:
        pdbqt_data = open(filename).readlines()

    pdb_block = ""
    # FIXME: Item "None" of "Optional[List[str]]" has no attribute "__iter__" (not iterable)
    for line in pdbqt_data:  # type: ignore
        pdb_block += "%s\n" % line[:66]
    return pdb_block


def convert_protein_to_pdbqt(mol: RDKitMol, outfile: str) -> None:
    """Convert a protein PDB file into a pdbqt file.

    Writes the extra PDBQT terms directly to `outfile`.

    Parameters
    ----------
    mol: rdkit.Chem.rdchem.Mol
        Protein molecule
    outfile: str
        filename which already has a valid pdb representation of mol
    """
    lines = [x.strip() for x in open(outfile).readlines()]
    out_lines = []
    for line in lines:
        if "ROOT" in line or "ENDROOT" in line or "TORSDOF" in line:
            out_lines.append("%s\n" % line)
            continue
        if not line.startswith("ATOM"):
            continue
        line = line[:66]
        atom_index = int(line[6:11])
        atom = mol.GetAtoms()[atom_index - 1]
        line = "%s    +0.000 %s\n" % (line, atom.GetSymbol().ljust(2))
        out_lines.append(line)
    with open(outfile, 'w') as fout:
        for line in out_lines:
            fout.write(line)


def _mol_to_graph(mol: RDKitMol):
    """Convert RDKit Mol to NetworkX graph

    Convert mol into a graph representation atoms are nodes, and bonds
    are vertices stored as graph

    Parameters
    ----------
    mol: rdkit.Chem.rdchem.Mol
        The molecule to convert into a graph.

    Returns
    -------
    graph: networkx.Graph
        Contains atoms indices as nodes, edges as bonds.

    Notes
    -----
    This function requires NetworkX to be installed.
    """
    try:
        import networkx as nx
    except ModuleNotFoundError:
        raise ImportError("This function requires NetworkX to be installed.")

    G = nx.Graph()
    num_atoms = mol.GetNumAtoms()
    G.add_nodes_from(range(num_atoms))
    for i in range(mol.GetNumBonds()):
        from_idx = mol.GetBonds()[i].GetBeginAtomIdx()
        to_idx = mol.GetBonds()[i].GetEndAtomIdx()
        G.add_edge(from_idx, to_idx)
    return G


def _get_rotatable_bonds(mol: RDKitMol) -> List[Tuple[int, int]]:
    """
    https://github.com/rdkit/rdkit/blob/f4529c910e546af590c56eba01f96e9015c269a6/Code/GraphMol/Descriptors/Lipinski.cpp#L107

    Taken from rdkit source to find which bonds are rotatable store
    rotatable bonds in (from_atom, to_atom)

    Parameters
    ----------
    mol: rdkit.Chem.rdchem.Mol
        Ligand molecule

    Returns
    -------
    rotatable_bonds: List[List[int, int]]
        List of rotatable bonds in molecule

    Notes
    -----
    This function requires RDKit to be installed.
    """
    try:
        from rdkit import Chem
        from rdkit.Chem import rdmolops
    except ModuleNotFoundError:
        raise ImportError("This function requires RDKit to be installed.")

    pattern = Chem.MolFromSmarts(
        "[!$(*#*)&!D1&!$(C(F)(F)F)&!$(C(Cl)(Cl)Cl)&!$(C(Br)(Br)Br)&!$(C([CH3])("
        "[CH3])[CH3])&!$([CD3](=[N,O,S])-!@[#7,O,S!D1])&!$([#7,O,S!D1]-!@[CD3]="
        "[N,O,S])&!$([CD3](=[N+])-!@[#7!D1])&!$([#7!D1]-!@[CD3]=[N+])]-!@[!$(*#"
        "*)&!D1&!$(C(F)(F)F)&!$(C(Cl)(Cl)Cl)&!$(C(Br)(Br)Br)&!$(C([CH3])([CH3])"
        "[CH3])]")
    rdmolops.FastFindRings(mol)
    rotatable_bonds = mol.GetSubstructMatches(pattern)
    return rotatable_bonds


def convert_mol_to_pdbqt(mol: RDKitMol, outfile: str) -> None:
    """Writes the provided ligand molecule to specified file in pdbqt format.

    Creates a torsion tree and write to pdbqt file. The torsion tree
    represents rotatable bonds in the molecule.

    Parameters
    ----------
    mol: rdkit.Chem.rdchem.Mol
        The molecule whose value is stored in pdb format in outfile
    outfile: str
        Filename for a valid pdb file with the extention .pdbqt

    Notes
    -----
    This function requires NetworkX to be installed.
    """
    try:
        import networkx as nx
    except ModuleNotFoundError:
        raise ImportError("This function requires NetworkX to be installed.")

    # Walk through the original file and extract ATOM/HETATM lines and
    # add PDBQT charge annotations.
    pdb_map = _create_pdb_map(outfile)
    graph = _mol_to_graph(mol)
    rotatable_bonds = _get_rotatable_bonds(mol)

    # Remove rotatable bonds from this molecule
    for bond in rotatable_bonds:
        graph.remove_edge(bond[0], bond[1])
    # Get the connected components now that the rotatable bonds have
    # been removed.
    components = [x for x in nx.connected_components(graph)]
    comp_map = _create_component_map(mol, components)

    used_partitions = set()
    lines = []
    # The root is the largest connected component.
    root = max(enumerate(components), key=lambda x: len(x[1]))[0]
    # Write the root component
    lines.append("ROOT\n")
    for atom in components[root]:
        lines.append(pdb_map[atom])
    lines.append("ENDROOT\n")
    # We've looked at the root, so take note of that
    used_partitions.add(root)
    for bond in rotatable_bonds:
        valid, next_partition = _valid_bond(used_partitions, bond, root,
                                            comp_map)
        if not valid:
            continue
        _dfs(used_partitions, next_partition, bond, components, rotatable_bonds,
             lines, pdb_map, comp_map)
    lines.append("TORSDOF %s" % len(rotatable_bonds))
    with open(outfile, 'w') as fout:
        for line in lines:
            fout.write(line)


def _create_pdb_map(outfile: str) -> Dict[int, str]:
    """Create a mapping from atom numbers to lines to write to pdbqt

    This is a map from rdkit atom number to its line in the pdb
    file. We also add the two additional columns required for
    pdbqt (charge, symbol).

    Note rdkit atoms are 0 indexed and pdb files are 1 indexed

    Parameters
    ----------
    outfile: str
        filename which already has a valid pdb representation of mol

    Returns
    -------
    pdb_map: Dict[int, str]
        Maps rdkit atom numbers to lines to be written to PDBQT file.
    """
    lines = [x.strip() for x in open(outfile).readlines()]
    lines = list(
        filter(lambda x: x.startswith("HETATM") or x.startswith("ATOM"), lines))
    lines = [x[:66] for x in lines]
    pdb_map = {}
    for line in lines:
        my_values = line.split()
        atom_number = int(my_values[1])
        atom_symbol = my_values[2]
        atom_symbol = ''.join([i for i in atom_symbol if not i.isdigit()])
        line = line.replace("HETATM", "ATOM  ")
        line = "%s    +0.000 %s\n" % (line, atom_symbol.ljust(2))
        pdb_map[atom_number - 1] = line
    return pdb_map


def _create_component_map(mol: RDKitMol,
                          components: List[List[int]]) -> Dict[int, int]:
    """Creates a map from atom ids to disconnected component id

    For each atom in `mol`, maps it to the id of the component in the
    molecule. The intent is that this is used on a molecule whose
    rotatable bonds have been removed. `components` is a list of the
    connected components after this surgery.

    Parameters
    ----------
    mol: rdkit.Chem.rdchem.Mol
        The molecule to find disconnected components in
    components: List[List[int]]
        List of connected components

    Returns
    -------
    comp_map: Dict[int, int]
        Maps atom ids to component ides
    """
    comp_map = {}
    for i in range(mol.GetNumAtoms()):
        for j in range(len(components)):
            if i in components[j]:
                comp_map[i] = j
                break
    return comp_map


def _dfs(used_partitions: Set[int], current_partition: int,
         bond: Tuple[int, int], components: List[List[int]],
         rotatable_bonds: List[Tuple[int, int]], lines: List[str],
         pdb_map: Dict[int, str], comp_map: Dict[int, int]) -> List[str]:
    """
    This function does a depth first search through the torsion tree

    Parameters
    ----------
    used_partions: Set[int]
        Partitions which have already been used
    current_partition: int
        The current partition to expand
    bond: Tuple[int, int]
        the bond which goes from the previous partition into this partition
    components: List[List[int]]
        List of connected components
    rotatable_bonds: List[Tuple[int, int]]
        List of rotatable bonds. This tuple is (from_atom, to_atom).
    lines: List[str]
        List of lines to write
    pdb_map: Dict[int, str]
        Maps atom numbers to PDBQT lines to write
    comp_map: Dict[int, int]
        Maps atom numbers to component numbers

    Returns
    -------
    lines: List[str]
        List of lines to write. This has more appended lines.
    """
    if comp_map[bond[1]] != current_partition:
        bond = (bond[1], bond[0])
    used_partitions.add(comp_map[bond[0]])
    used_partitions.add(comp_map[bond[1]])
    lines.append("BRANCH %4s %4s\n" % (bond[0] + 1, bond[1] + 1))
    for atom in components[current_partition]:
        lines.append(pdb_map[atom])
    for b in rotatable_bonds:
        valid, next_partition = \
          _valid_bond(used_partitions, b, current_partition, comp_map)
        if not valid:
            continue
        lines = _dfs(used_partitions, next_partition, b, components,
                     rotatable_bonds, lines, pdb_map, comp_map)
    lines.append("ENDBRANCH %4s %4s\n" % (bond[0] + 1, bond[1] + 1))
    return lines


def _valid_bond(used_partitions: Set[int], bond: Tuple[int, int],
                current_partition: int,
                comp_map: Dict[int, int]) -> Tuple[bool, int]:
    """Helper method to find next partition to explore.

    Used to check if a bond goes from the current partition into a
    partition that is not yet explored

    Parameters
    ----------
    used_partions: Set[int]
        Partitions which have already been used
    bond: Tuple[int, int]
        The bond to check if it goes to an unexplored partition.
        This tuple is (from_atom, to_atom).
    current_partition: int
        The current partition of the DFS
    comp_map: Dict[int, int]
        Maps atom ids to component ids

    Returns
    -------
    is_valid: bool
        Whether to exist the next partition or not
    next_partition: int
        The next partition to explore
    """
    part1 = comp_map[bond[0]]
    part2 = comp_map[bond[1]]
    if part1 != current_partition and part2 != current_partition:
        return False, 0
    if part1 == current_partition:
        next_partition = part2
    else:
        next_partition = part1
    return next_partition not in used_partitions, next_partition
