# Copyright 2026 AlQuraishi Laboratory
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import subprocess as sp
import tempfile
from collections import defaultdict
from pathlib import Path
from typing import Literal

import pandas as pd
from rdkit import Chem
from rdkit.Chem import rdFingerprintGenerator
from tqdm import tqdm

from openfold3.core.data.io.sequence.fasta import write_multichain_fasta
from openfold3.core.data.primitives.caches.filtering import (
    get_all_cache_chains,
    get_mol_id_to_smiles,
    logger,
)
from openfold3.core.data.primitives.caches.format import (
    ClusteredDatasetCache,
    ValidationDatasetCache,
)
from openfold3.core.data.resources.residues import MoleculeType


def precompute_fingerprints(
    smiles_list: list[str], mfpgen: rdFingerprintGenerator
) -> list[Chem.DataStructs.ExplicitBitVect | None]:
    """Precompute fingerprints for a list of SMILES strings.

    Args:
        smiles_list (list[str]): List of SMILES strings.
        mfpgen (rdFingerprintGenerator): RDKit fingerprint generator.

    Returns:
        list[Chem.DataStructs.ExplicitBitVect | None]:
            List of fingerprints. If a fingerprint could not be generated, None is
            returned.

    """
    mols = [Chem.MolFromSmiles(smiles) for smiles in smiles_list]

    fingerprints = []

    for index, mol in enumerate(mols):
        if mol is None:
            logger.warning(
                f"Error in generating fingerprint for molecule: {smiles_list[index]}"
            )
            fingerprints.append(None)
        else:
            fingerprints.append(mfpgen.GetFingerprint(mol))

    return fingerprints


def get_mol_id_to_tanimoto_ligands(
    val_dataset_cache: ClusteredDatasetCache,
    train_dataset_cache: ClusteredDatasetCache,
    similarity_threshold: float = 0.85,
) -> dict[str, set[str]]:
    """
    Identify ligands in the validation dataset that have high Tanimoto similarity (>
    similarity_threshold) with any ligand in the training dataset.

    Args:
        val_dataset_cache (ClusteredDatasetCache):
            Validation dataset cache.
        train_dataset_cache (ClusteredDatasetCache):
            Training dataset cache.
        similarity_threshold (float):
            Tanimoto similarity above which a ligand is considered homologous.

    Returns:
        Dict[str, set[str]]:
            Mapping of ligand reference IDs in the validation set to homologous ligand
            reference IDs in the training set.
    """
    logger.info("Computing fingerprints.")

    # Initialize fingerprint generator once
    mfpgen = rdFingerprintGenerator.GetMorganGenerator(radius=3, fpSize=2048)

    # Extract SMILES and reference IDs from validation dataset
    val_refid_to_smiles = get_mol_id_to_smiles(val_dataset_cache)

    # Extract SMILES and reference IDs from training dataset
    train_refid_to_smiles = get_mol_id_to_smiles(train_dataset_cache)

    # Precompute fingerprints for validation SMILES
    val_refids = list(val_refid_to_smiles.keys())
    val_smiles = list(val_refid_to_smiles.values())
    val_fps = precompute_fingerprints(val_smiles, mfpgen)

    assert len(val_refids) == len(val_fps)
    val_refid_to_fp = dict(zip(val_refids, val_fps, strict=True))

    # Precompute fingerprints for training SMILES
    train_refids = list(train_refid_to_smiles.keys())
    train_smiles = list(train_refid_to_smiles.values())
    train_fps = precompute_fingerprints(train_smiles, mfpgen)

    assert len(train_refids) == len(train_fps)

    if None in train_fps:
        raise ValueError(
            "Some training-set molecule fingerprints are None. Can't reliably calculate"
            " training-set homology."
        )

    # Make map of ligands to high-homology-related ligands
    val_refid_to_homologs = defaultdict(set)

    # Iterate over validation set ligands
    logger.info("Computing ligand similarities.")
    for val_refid, val_fp in tqdm(val_refid_to_fp.items()):
        # If fingerprint is not available, conservatively set homologous ligands
        # to everything
        if val_fp is None:
            val_refid_to_homologs[val_refid] = set(train_refids)
            continue

        # Compute similarities in bulk
        similarities = Chem.DataStructs.BulkTanimotoSimilarity(val_fp, train_fps)

        # Get all similar ligands in training set
        homolog_train_refids = {
            train_refids[i]
            for i, score in enumerate(similarities)
            if score > similarity_threshold
        }

        # Set homologous ligands for this ligand
        val_refid_to_homologs[val_refid] = homolog_train_refids

    assert set(val_refid_to_homologs.keys()) == set(val_refids)

    return val_refid_to_homologs


# TODO: Add back thread and memory settings if they keep being a problem with newer
# MMSeqs versions
def run_mmseqs_search(
    query_id_to_sequence: dict[str, str],
    target_id_to_sequence: dict[str, str],
    seq_identity_threshold: float = 0.4,
    seq_identity_mode: Literal["global", "local"] = "global",
    sensitivity: float = 8.0,
    mmseqs_binary: str = "mmseqs",
    dbtype: int = 0,
) -> dict[str, set[str]]:
    """
    Perform MMseqs2 search between query and target sequences and return pairs above a
    certain sequence identity.

    Args:
        query_id_to_sequence (Dict[str, str]):
            Mapping of query sequence IDs to sequences.
        target_id_to_sequence (Dict[str, str]):
            Mapping of target sequence IDs to sequences.
        seq_identity_threshold (float):
            Sequence identity threshold above which a hit is considered homologous. Note
            that this refers to the global sequence identity with respect to the full
            length of the query sequence.
        sensitivity (float):
            Sensitivity of the search.
        mmseqs_binary (str):
            Path to the MMseqs2 binary.
        dbtype (int):
            Database type passed to MMSeqs2. 0 for auto, 1 for amino acids, 2 for
            nucleotides.
    Returns:
        Dict[str, set[str]]:
            Mapping of query sequence IDs to homologous target sequence IDs.
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir = Path(temp_dir)
        query_fasta = temp_dir / "query.fasta"
        target_fasta = temp_dir / "target.fasta"
        db_target = temp_dir / "target_db"
        db_query = temp_dir / "query_db"
        db_result = temp_dir / "result_db"
        result_tsv = temp_dir / "search_result.tsv"

        # Write query and target sequences to FASTA files
        write_multichain_fasta(query_fasta, query_id_to_sequence)
        write_multichain_fasta(target_fasta, target_id_to_sequence)

        # Create MMseqs2 database for target
        cmd_make_db = (
            f"{mmseqs_binary} createdb {target_fasta} {db_target} --dbtype {dbtype}"
        )
        logger.info("Creating MMseqs2 database for target sequences.")
        sp.run(cmd_make_db, shell=True, check=True)

        cmd_make_db = (
            f"{mmseqs_binary} createdb {query_fasta} {db_query} --dbtype {dbtype}"
        )
        logger.info("Creating MMseqs2 database for query sequences.")
        sp.run(cmd_make_db, shell=True, check=True)

        # Decide search setting by dbtype
        if dbtype == 1:
            search_type = 1  # amino acid
            num_iterations = 3  # use profile search for proteins
        elif dbtype == 2:
            search_type = 3  # nucleotide
            num_iterations = 1
        else:
            search_type = 0  # auto
            num_iterations = 1

        # Run MMseqs2 search with high sensitivity and ensuring that all target
        # sequences can be included in the hits. Also prefilter by coverage wrt the
        # query sequence as that is guaranteed to be >= query sequence identity.
        cmd_search = (
            f"{mmseqs_binary} search "
            f"{db_query} "
            f"{db_target} "
            f"{db_result} "
            f"{temp_dir}/tmp "
            f"-s {sensitivity} "
            f"--max-seqs {len(target_id_to_sequence)} "
            f"--search-type {search_type} "
            f"-c {seq_identity_threshold} "
            "--cov-mode 2 "
            f"--num-iterations {num_iterations} "
            "-a "
            "-e 100"  # rely mostly on seq identity and coverage filters
        )
        logger.info("Running MMseqs2 search.")
        sp.run(cmd_search, shell=True, check=True)

        # Convert results to tabular format
        cmd_convert = (
            f"{mmseqs_binary} convertalis "
            f"{db_query} "
            f"{db_target} "
            f"{db_result} "
            f"{result_tsv} "
            "--format-output 'query,target,fident,nident,qlen'"
        )
        logger.info("Converting MMseqs2 search results to TSV.")
        sp.run(cmd_convert, shell=True, check=True)

        # Parse the search results
        logger.info("Parsing MMseqs2 search results.")
        df = pd.read_csv(result_tsv, sep="\t", header=None)

        if df.empty:
            logger.warning("MMseqs2 search returned no hits.")
            return {}

        assert df[2].min() >= 0.0 and df[2].max() <= 1.0  # assert column parsing

        logger.info("Filtering templates.")
        query_seq_id_to_homologs = {}

        if seq_identity_mode == "global":
            # Global query sequence identity = nident / qlen
            query_sequence_identity = df[3] / df[4]
        elif seq_identity_mode == "local":
            query_sequence_identity = df[2]
        else:
            raise ValueError(
                f"seq_identity_mode must be 'global' or 'local', got "
                f"{seq_identity_mode}"
            )

        high_identity = df[query_sequence_identity > seq_identity_threshold]

        # Group by column 0 (query ID), then collect sets of column 1 (homolog IDs)
        grouped = high_identity.groupby(0)[1].apply(set)

        for query_id, homologs in tqdm(grouped.items()):
            query_seq_id_to_homologs[query_id] = homologs

    logger.info(f"Found hits for {len(query_seq_id_to_homologs)} sequences.")

    return query_seq_id_to_homologs


def get_polymer_chain_to_homolog_chains(
    val_dataset_cache: ClusteredDatasetCache,
    train_dataset_cache: ClusteredDatasetCache,
    val_id_to_sequence: dict[str, str],
    train_id_to_sequence: dict[str, str],
    seq_identity_threshold: float = 0.4,
) -> dict[str, set[str]]:
    """Maps protein/nucleic-acid validation chains to homologous training chains.

    This uses MMSeqs template search to find homologous chains between the validation
    and training datasets. The search is run separately for protein and nucleotide
    chains to prevent any sequence ambiguity.

    Args:
        val_dataset_cache (ClusteredDatasetCache):
            Validation dataset cache.
        train_dataset_cache (ClusteredDatasetCache):
            Training dataset cache.
        val_id_to_sequence (Dict[str, str]):
            Mapping of {PDB-ID_chain-ID: sequence} for chains in the validation cache.
        train_id_to_sequence (Dict[str, str]):
            Mapping of {PDB-ID_chain-ID: sequence} for chains in the training cache.
        seq_identity_threshold (float):
            Sequence identity threshold above which a hit is considered homologous. Note
            that this refers to the global sequence identity with respect to the full
            length of the query sequence.

    Returns:
        Dict[str, set[str]]:
            Mapping of validation {pdb_id}_{chain_id} to homologous training
            {pdb_id}_{chain_id}s.
    """
    val_structure_cache = val_dataset_cache.structure_data
    train_structure_cache = train_dataset_cache.structure_data

    # Dictionary keyed on validation {pdb_id}_{chain_id} mapping to the homologous
    # training {pdb_id}_{chain_id}s
    chain_to_homologs = {}

    # Run MMSeqs search separately for protein and nucleotide chains to prevent any
    # sequence ambiguity
    for molecule_types in (
        (MoleculeType.PROTEIN,),
        (MoleculeType.DNA, MoleculeType.RNA),
    ):
        # Only get chains that have that molecule type
        val_chains = get_all_cache_chains(
            val_structure_cache, restrict_to_molecule_types=molecule_types
        )
        train_chains = get_all_cache_chains(
            train_structure_cache, restrict_to_molecule_types=molecule_types
        )

        val_id_to_sequence_mol = {k: val_id_to_sequence[k] for k in val_chains}
        train_id_to_sequence_mol = {k: train_id_to_sequence[k] for k in train_chains}

        if molecule_types == (MoleculeType.PROTEIN,):
            logger.info("Running MMseqs2 search for protein chains.")
            db_type = 1
        else:
            logger.info("Running MMseqs2 search for nucleotide chains.")
            db_type = 2

        chain_to_homologs.update(
            run_mmseqs_search(
                query_id_to_sequence=val_id_to_sequence_mol,
                target_id_to_sequence=train_id_to_sequence_mol,
                seq_identity_threshold=seq_identity_threshold,
                dbtype=db_type,
            )
        )

    return chain_to_homologs


def assign_homology_labels(
    val_dataset_cache: ValidationDatasetCache,
    train_dataset_cache: ClusteredDatasetCache,
    val_id_to_sequence: dict[str, str],
    train_id_to_sequence: dict[str, str],
    seq_identity_threshold: float = 0.4,
    tanimoto_threshold: float = 0.85,
) -> ValidationDatasetCache:
    """Detects if chains/interfaces are low-homology to the training dataset.

    Following AF3 SI 5.8, this function labels chains as low homology if there is no
    chain in the training set with a sequence identity greater than the threshold for
    polymer chains, or a Tanimoto similarity greater than the threshold for ligands.

    Interfaces in the validation dataset are labeled as low homology if there is no PDB
    in the training set that contains homologous chains to both chains in the interface.

    Args:
        val_dataset_cache (ValClusteredDatasetCache):
            Validation dataset cache.
        train_dataset_cache (ClusteredDatasetCache):
            Training dataset cache.
        val_id_to_sequence (Dict[str, str]):
            Mapping of {PDB-ID_chain-ID: sequence} for chains in the validation cache.
        train_id_to_sequence (Dict[str, str]):
            Mapping of {PDB-ID_chain-ID: sequence} for chains in the training cache.
        seq_identity_threshold (float):
            Sequence identity threshold. Chains with sequence identity strictly greater
            than this value are considered homologous.
        tanimoto_threshold (float):
            Tanimoto similarity threshold. Ligands with Tanimoto similarity strictly
            greater than this value are considered homologous.
    """
    val_structure_cache = val_dataset_cache.structure_data
    train_structure_cache = train_dataset_cache.structure_data

    # PREPARE LIGAND SIMILARITY
    # Get ligands with high Tanimoto similarity
    val_mol_id_to_homologs = get_mol_id_to_tanimoto_ligands(
        val_dataset_cache, train_dataset_cache, similarity_threshold=tanimoto_threshold
    )

    # Get mapping of every training set reference molecule ID to all PDB-IDs it occurs
    # in.
    train_mol_id_to_pdb_ids = defaultdict(set)
    for pdb_id, metadata in train_structure_cache.items():
        for chain_metadata in metadata.chains.values():
            if chain_metadata.molecule_type == MoleculeType.LIGAND:
                train_mol_id_to_pdb_ids[chain_metadata.reference_mol_id].add(pdb_id)

    # Map validation molecule IDs to training PDB IDs containing a homologous ligand
    val_mol_id_to_homolog_pdbs = {
        val_mol_id: set.union(
            *(train_mol_id_to_pdb_ids[train_mol_id] for train_mol_id in train_mol_ids)
        )
        if train_mol_ids
        else set()
        for val_mol_id, train_mol_ids in val_mol_id_to_homologs.items()
    }

    # PREPARE SEQUENCE SIMILARITY
    val_chain_to_homologs = get_polymer_chain_to_homolog_chains(
        val_dataset_cache=val_dataset_cache,
        train_dataset_cache=train_dataset_cache,
        val_id_to_sequence=val_id_to_sequence,
        train_id_to_sequence=train_id_to_sequence,
        seq_identity_threshold=seq_identity_threshold,
    )

    # Similarly to earlier, map validation chain IDs to training PDB IDs containing a
    # homologous chain
    val_chain_to_homolog_pdbs = {
        val_chain: set(pdb_chain_id[:4] for pdb_chain_id in train_chain_ids)
        for val_chain, train_chain_ids in val_chain_to_homologs.items()
    }

    def get_homolog_pdbs(pdb_id: str, chain_id: str, chain_data):
        """Small helper function to retrieve the set of homolog PDBs for a chain."""
        pdb_chain_id = f"{pdb_id}_{chain_id}"

        # Get homologous PDBs for ligand chains
        if chain_data.molecule_type == MoleculeType.LIGAND:
            return val_mol_id_to_homolog_pdbs.get(chain_data.reference_mol_id, set())
        # Get homologous PDBs for polymer chains
        else:
            return val_chain_to_homolog_pdbs.get(pdb_chain_id, set())

    # SET HOMOLOGY LABELS
    # Assign homology labels to validation dataset
    logger.info("Setting homology labels.")
    for pdb_id, structure_data in tqdm(val_structure_cache.items()):
        # Start by setting chain-wise homology
        for chain_id, chain_data in structure_data.chains.items():
            # Set homology to low if there is no training PDB with a homologous chain
            homolog_pdbs = get_homolog_pdbs(pdb_id, chain_id, chain_data)
            chain_data.low_homology = len(homolog_pdbs) == 0

        # Continue with interface-wise homology
        for interface_id, interface_data in structure_data.interfaces.items():
            chain_id_1, chain_id_2 = interface_id.split("_")
            chain_data_1 = structure_data.chains[chain_id_1]
            chain_data_2 = structure_data.chains[chain_id_2]

            pdbs_with_homolog_chain_1 = get_homolog_pdbs(
                pdb_id, chain_id_1, chain_data_1
            )
            pdbs_with_homolog_chain_2 = get_homolog_pdbs(
                pdb_id, chain_id_2, chain_data_2
            )

            # Set to low homology if there is no single PDB-ID containing chains
            # homologous to both interface chains
            interface_data.low_homology = pdbs_with_homolog_chain_1.isdisjoint(
                pdbs_with_homolog_chain_2
            )

    return val_dataset_cache
