import numpy as np


def pairwise_distances(A, B):
    diff = A[:, None, :] - B[None, :, :]
    return np.linalg.norm(diff, axis=-1)


def lennard_jones_repulsion(dist, sigma=3.0, epsilon=0.2, cutoff=8.0):
    """
    Simple clash/repulsion proxy.
    Returns only the repulsive part.
    """
    mask = dist < cutoff
    d = np.clip(dist[mask], 1e-6, None)
    sr6 = (sigma / d) ** 6
    rep = 4 * epsilon * sr6**2
    return rep.sum()


def count_contacts(dist, cutoff):
    return int(np.sum(dist < cutoff))


def count_hbonds(
    ligand_pos,
    ligand_types,
    protein_pos,
    protein_types,
    cutoff=3.5,
):
    """
    Simplified H-bond count using atom labels only.
    Better version should include angular criteria and donor/acceptor typing.
    """
    donor_acceptor = {"N", "O", "S", "F"}
    lig_mask = np.array([t in donor_acceptor for t in ligand_types])
    pro_mask = np.array([t in donor_acceptor for t in protein_types])

    if not lig_mask.any() or not pro_mask.any():
        return 0

    d = pairwise_distances(ligand_pos[lig_mask], protein_pos[pro_mask])
    return int(np.sum(d < cutoff))


def ligand_strain_proxy(
    ligand_pos,
    reference_ligand_pos,
):
    """
    Very cheap internal deformation proxy.
    """
    return float(np.mean(np.sum((ligand_pos - reference_ligand_pos) ** 2, axis=1)))


def buried_surface_proxy(ligand_pos, protein_pos, contact_cutoff=4.5):
    d = pairwise_distances(ligand_pos, protein_pos)
    n_contacts = np.sum(d < contact_cutoff)
    return float(n_contacts / max(len(ligand_pos), 1))


def score_structure(
    atoms,
    ligand_idx,
    protein_idx,
    ligand_types,
    protein_types,
    reference_ligand_pos=None,
    mm_interaction_energy=None,
    mm_ligand_strain=None,
) -> np.ndarray:
    """
    Returns a scoring vector for multi-objective optimization.

    Indices:
    0: clash_score (minimize)
    1: ligand strain (minimize)
    2: -hbonds (minimize, means maximize)
    3: -hydrophobic_contacts (minimize, means maximize)
    4: -burial (minimize, means maximize)
    5: MM Interaction Energy (optional)
    6: MM Ligand Strain Energy (optional)
    """
    pos = atoms.get_positions()

    lig_pos = pos[ligand_idx]
    pro_pos = pos[protein_idx]

    d_lp = pairwise_distances(lig_pos, pro_pos)

    clash_score = lennard_jones_repulsion(d_lp, sigma=2.8, epsilon=0.1, cutoff=6.0)
    hydrophobic_contacts = count_contacts(d_lp, cutoff=4.2)
    hbonds = count_hbonds(lig_pos, ligand_types, pro_pos, protein_types, cutoff=3.5)
    burial = buried_surface_proxy(lig_pos, pro_pos, contact_cutoff=4.5)

    if reference_ligand_pos is None:
        strain = 0.0
    else:
        strain = ligand_strain_proxy(lig_pos, reference_ligand_pos)

    vec = [
        clash_score,             # minimize
        strain,                  # minimize
        -hbonds,                 # minimize
        -hydrophobic_contacts,   # minimize
        -burial,                 # minimize
    ]

    if mm_interaction_energy is None:
        pass
    else:
        vec.append(mm_interaction_energy)
        if mm_ligand_strain is not None:
            vec.append(mm_ligand_strain)

    return np.array(vec, dtype=float)
