# Copyright 2025 ByteDance and/or its affiliates.
#
# 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.

from collections import Counter
from datetime import datetime

import biotite.sequence as seq
import biotite.sequence.align as align
import networkx as nx
import numpy as np
from biotite.structure import Atom, AtomArray, get_chain_starts, get_residue_starts
from rdkit import Chem

from pxmeter.constants import PRO_STD_RESIDUES_ONE_LETTER


def get_unique_chain_id(atom_array: AtomArray) -> np.ndarray:
    """
    Generate unique chain IDs for each chain in the AtomArray.

    This function takes an AtomArray and assigns a unique string ID to each chain.
    The ID is based on the original chain ID, with a numerical suffix added to ensure uniqueness.
    Example: [A, B, A, B, C] -> [A, B, A.1, B.1, C]

    Args:
        atom_array (AtomArray): The input AtomArray containing the atomic coordinates and metadata.

    Returns:
        np.ndarray: An array of unique chain IDs, with the same length as the input AtomArray.
    """
    chain_ids = np.zeros(len(atom_array), dtype="<U8")  # <U8 for very large structure
    chain_starts = get_chain_starts(atom_array, add_exclusive_stop=True)

    chain_counter = Counter()
    for start, stop in zip(chain_starts[:-1], chain_starts[1:]):
        ori_chain_id = atom_array.chain_id[start]
        cnt = chain_counter[ori_chain_id]
        if cnt == 0:
            new_chain_id = ori_chain_id
        else:
            new_chain_id = f"{ori_chain_id}.{chain_counter[ori_chain_id]}"

        chain_ids[start:stop] = new_chain_id
        chain_counter[ori_chain_id] += 1

    assert "" not in chain_ids
    return chain_ids


def get_unique_atom_id(atom_array: AtomArray) -> np.ndarray:
    """
    Generate unique atom identifiers by combining residue IDs, residue names, and atom names.

    Args:
        res_id (np.ndarray): An array of residue IDs.
        res_name (np.ndarray): An array of residue names.
        atom_name (np.ndarray): An array of atom names.

    Returns:
        np.ndarray: An array of unique atom identifiers, each formed by concatenating
                    the corresponding residue ID, residue name, and atom name with underscores.
    """
    stacked = np.column_stack(
        (atom_array.res_id, atom_array.res_name, atom_array.atom_name)
    )
    unique_atom_id = np.array(["_".join(map(str, row)) for row in stacked])
    return unique_atom_id


def get_chain_ccd_seq(atom_array: AtomArray) -> dict[str, str]:
    """
    Generate a dictionary mapping chain IDs to CCD sequences.

    This function takes an AtomArray and extracts the CCD sequences for each chain.
    The CCD sequence is a string representation of the residue names in the chain,
    separated by underscores.

    Args:
        atom_array (AtomArray): The input AtomArray containing the atomic coordinates and metadata.

    Returns:
        dict[str, str]: A dictionary where the keys are chain IDs and the values are CCD sequences.
    """
    chain_id_to_ccd_seq = {}
    for chain_id in np.unique(atom_array.chain_id):
        chain_array = atom_array[atom_array.chain_id == chain_id]
        res_starts = get_residue_starts(chain_array)
        ccd_array = chain_array.res_name[res_starts]
        ccd_seq = "_".join(ccd_array)
        chain_id_to_ccd_seq[chain_id] = ccd_seq
    return chain_id_to_ccd_seq


def get_seq_alignment_score(seq1: str, seq2: str, seq_type: str = "protein") -> tuple:
    """
    Calculate the alignment score between two sequences.
    The global alignment based on the Needleman-Wunsch algorithm.

    Args:
        seq1 (str): The first sequence string.
        seq2 (str): The second sequence string.
        seq_type (str): The type of sequences to align.
            One of {"protein", "nuc"} for protein or nucleotide sequences. Default is "protein".

    Returns:
        tuple: A tuple containing:
            - alignments (list): A list of optimal alignments.
            - sequence_identity (float): The sequence identity of the first optimal alignment.

    Raises:
        NotImplementedError: If the seq_type is not "protein" or "nuc".
    """

    def _std_prot_seq(prot_seq):
        std_seq = ""
        for i in prot_seq:
            if i in PRO_STD_RESIDUES_ONE_LETTER:
                std_seq += i
            elif i == "U":
                # U -> SEC -> CYS -> C
                std_seq += "C"
            else:
                std_seq += "X"
        return std_seq

    def _std_nuc_seq(nuc_seq):
        std_seq = ""
        for i in nuc_seq:
            if i in {"A", "T", "C", "G"}:
                std_seq += i
            elif i == "U":
                std_seq += "T"
            else:
                std_seq += "N"
        return std_seq

    if seq_type == "protein":
        seq1 = seq.ProteinSequence(_std_prot_seq(seq1))
        seq2 = seq.ProteinSequence(_std_prot_seq(seq2))
        matrix = align.SubstitutionMatrix.std_protein_matrix()
    elif seq_type == "nuc":
        seq1 = seq.NucleotideSequence(_std_nuc_seq(seq1))
        seq2 = seq.NucleotideSequence(_std_nuc_seq(seq2))
        matrix = align.SubstitutionMatrix.std_nucleotide_matrix()
    else:
        raise NotImplementedError(f"seq_type {seq_type} is not supported")

    alignments = align.align_optimal(seq1, seq2, matrix, gap_penalty=-10, local=False)
    return alignments[0], align.get_sequence_identity(alignments[0])


def get_inter_residue_bonds(atom_array: AtomArray) -> np.ndarray:
    """
    Get inter residue bonds by checking chain_id and res_id

    Args:
        atom_array (AtomArray): Biotite AtomArray, must have chain_id and res_id

    Returns:
        np.ndarray: inter residue bonds, shape = (n,2)
    """
    if atom_array.bonds is None:
        return np.array([])
    idx_i = atom_array.bonds._bonds[:, 0]
    idx_j = atom_array.bonds._bonds[:, 1]
    chain_id_diff = atom_array.chain_id[idx_i] != atom_array.chain_id[idx_j]
    res_id_diff = atom_array.res_id[idx_i] != atom_array.res_id[idx_j]
    diff_mask = chain_id_diff | res_id_diff
    inter_residue_bonds = atom_array.bonds._bonds[diff_mask]
    inter_residue_bonds = inter_residue_bonds[:, :2]  # remove bond type
    return inter_residue_bonds


def create_single_atom_array(atom: Atom) -> AtomArray:
    """
    Create a single atom AtomArray from a single atom.
    The biotite.structure.array function would change the dtype of the annotation categories.
    For example, if the Atom has annotation category "label_entity_id", the dtype of Atom is "<U2",
    but the dtype of AtomArray is "<U1".

    Args:
        atom (Atom): Biotite Atom object

    Returns:
        AtomArray: A AtomArray object with a single atom.
    """
    array = AtomArray(1)
    # Add all (also optional) annotation categories
    for name in atom._annot.keys():
        if name in array._annot:
            array.del_annotation(name)
        array.set_annotation(name, [atom._annot[name]])

    # Add all atoms to AtomArray
    array._coord[0] = atom.coord
    return array


def rdkit_mol_to_nx_graph(mol: Chem.Mol) -> nx.Graph:
    """
    Convert an RDKit molecule to a NetworkX graph representation.

    Args:
        mol (Chem.Mol): Input RDKit molecule object

    Returns:
        nx.Graph: NetworkX graph where nodes represent atoms and edges represent bonds
    """
    G = nx.Graph()

    for atom in mol.GetAtoms():
        G.add_node(
            atom.GetIdx(),
            atomic_num=atom.GetAtomicNum(),
            is_aromatic=atom.GetIsAromatic(),
            atom_symbol=atom.GetSymbol(),
        )

    for bond in mol.GetBonds():
        G.add_edge(
            bond.GetBeginAtomIdx(), bond.GetEndAtomIdx(), bond_type=bond.GetBondType()
        )

    return G


def get_mol_graph_matches(
    mol_graph1: nx.Graph, mol_graph2: nx.Graph, max_matches: int = 1000
) -> list[dict]:
    """
    Find all isomorphisms between subgraph of mol_graph1 and mol_graph2.

    Args:
        mol_graph1 (nx.Graph): Source molecular graph (typically larger subgraph)
        mol_graph2 (nx.Graph): Target molecular graph (typically smaller supergraph)
        max_matches (int): Maximum number of matches to return. Default is 1000.

    Returns:
        list[dict]: List of node mapping dictionaries where each dictionary represents
        a subgraph isomorphism (mapping from mol_graph1 node indices to mol_graph2 node indices)

    Note:
        Uses atomic number comparison for node matching, ignoring other atom properties
    """
    isomatcher = nx.algorithms.isomorphism.GraphMatcher(
        mol_graph1,
        mol_graph2,
        node_match=lambda x, y: x["atomic_num"] == y["atomic_num"],
    )

    matches = []
    num = 0
    for i in isomatcher.subgraph_isomorphisms_iter():
        matches.append(i)
        num += 1
        if num >= max_matches:
            break
    return matches


def get_res_graph_matches(
    res_graph1: nx.Graph, res_graph2: nx.Graph, max_matches: int = 1000
) -> list[dict]:
    """
    Find subgraph isomorphisms between two residue-level graphs using residue names.

    This function enumerates mappings where a subgraph of `res_graph1` is isomorphic to
    (i.e., can be relabeled to match) `res_graph2`. Node equivalence is determined
    solely by the `"res_name"` node attribute; all other node or edge attributes are ignored.
    Enumeration stops once `max_matches` mappings have been collected.

    Args:
        res_graph1 (nx.Graph): The source (typically larger) residue graph.
            Node attribute required: ``"res_name"`` (e.g., "ALA", "NAG").
        res_graph2 (nx.Graph): The target (typically smaller) residue graph to match against.
            Node attribute required: ``"res_name"``.
        max_matches (int, optional): Maximum number of mappings to return. Defaults to ``1000``.

    Returns:
        list[dict]: A list of node-mapping dicts. Each dict maps node IDs from `res_graph1`
        (keys) to node IDs in `res_graph2` (values) representing one subgraph isomorphism.
    """
    isomatcher = nx.algorithms.isomorphism.GraphMatcher(
        res_graph1,
        res_graph2,
        node_match=lambda x, y: (x["res_name"] == y["res_name"])
        and (x["atom_names"] == y["atom_names"]),
    )

    matches = []
    num = 0
    for i in isomatcher.subgraph_isomorphisms_iter():
        matches.append(i)
        num += 1
        if num >= max_matches:
            break
    return matches


def is_valid_date_format(date_string: str) -> bool:
    """
    Check if the date string is in the format yyyy-mm-dd.

    Args:
        date_string (str): The date string to check.

    Returns:
        bool: True if the date string is in the format yyyy-mm-dd, False otherwise.
    """
    try:
        datetime.strptime(date_string, "%Y-%m-%d")
        return True
    except ValueError:
        return False
