# Copyright 2024 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.

import torch

from protenix.metrics.rmsd import rmsd
from protenix.model.utils import expand_at_dim
from protenix.utils.distributed import traverse_and_aggregate
from protenix.utils.logger import get_logger
from protenix.utils.permutation.chain_permutation.utils import (
    apply_transform,
    get_optimal_transform,
    num_unique_matches,
)

logger = get_logger(__name__)


def permute_pred_to_optimize_pocket_aligned_rmsd(
    pred_coord: torch.Tensor,  # [N_sample, N_atom, 3]
    true_coord: torch.Tensor,  # [N_atom, 3]
    true_coord_mask: torch.Tensor,
    true_pocket_mask: torch.Tensor,
    true_ligand_mask: torch.Tensor,
    atom_entity_id: torch.Tensor,  # [N_atom]
    atom_asym_id: torch.Tensor,  # [N_atom]
    mol_atom_index: torch.Tensor,  # [N_atom]
    use_center_rmsd: bool = False,
):
    """

    Returns:
        permute_pred_indices (list[torch.Tensor]): A list of LongTensor.
            The list contains N_sample elements.
            Each elements is a LongTensor of shape = [N_atom].
        permuted_aligned_pred_coord (torch.Tensor): permuted and aligned coordinates of pred_coord.
            [N_sample, N_atom, 3]
    """

    log_dict = {}
    atom_entity_id = atom_entity_id.long()
    atom_asym_id = atom_asym_id.long()
    mol_atom_index = mol_atom_index.long()
    true_coord_mask = true_coord_mask.bool()
    true_pocket_mask = true_pocket_mask.bool()
    true_ligand_mask = true_ligand_mask.bool()
    assert pred_coord.size(-2) == true_coord.size(-2), "Atom numbers are difference."
    assert pred_coord.dim() == 3

    # find entity_id/asym_id of pocket and ligand chains
    def _get_entity_and_asym_id(atom_mask):

        masked_asym_id = atom_asym_id[atom_mask]
        masked_entity_id = atom_entity_id[atom_mask]
        assert (masked_asym_id[0] == masked_asym_id).all()
        assert (masked_entity_id[0] == masked_entity_id).all()
        return masked_asym_id[0].item(), masked_entity_id[0].item()

    pocket_asym_id, pocket_entity_id = _get_entity_and_asym_id(true_pocket_mask)
    ligand_asym_id, ligand_entity_id = _get_entity_and_asym_id(true_ligand_mask)

    candidate_pockets = {}
    for i in torch.unique(atom_asym_id[atom_entity_id == pocket_entity_id]):
        i = i.item()
        pocket_mask = atom_asym_id == i
        pocket_mask = pocket_mask * torch.isin(
            mol_atom_index, mol_atom_index[true_pocket_mask]
        )
        assert pocket_mask.sum() == true_pocket_mask.sum()
        candidate_pockets[i] = pocket_mask.clone()
    candidate_ligands = {}
    for j in torch.unique(atom_asym_id[atom_entity_id == ligand_entity_id]):
        j = j.item()
        lig_mask_j = atom_asym_id == j
        if lig_mask_j.sum() != true_ligand_mask.sum():
            logger.warning(
                f"The ligand selected by 'mol_id' has {lig_mask_j.sum().item()} atoms."
                + f"The true ligand selected by 'asym_id' has {true_ligand_mask.sum().item()} atoms."
            )
            lig_mask_j = lig_mask_j * torch.isin(
                mol_atom_index, mol_atom_index[true_ligand_mask]
            )
        assert lig_mask_j.sum() == true_ligand_mask.sum()
        candidate_ligands[j] = lig_mask_j

    log_dict["num_sym_pocket"] = len(candidate_pockets)
    log_dict["num_sym_ligand"] = len(candidate_ligands)
    log_dict["has_sym_chain"] = len(candidate_ligands) + len(candidate_pockets) > 2

    # Enumerate over the batch dimension of pred_coord
    # to find the optimal chain assignment for each sample.

    def _find_protein_ligand_chains_for_one_sample(
        coord: torch.Tensor,
    ):
        best_results = {}
        unpermuted_results = {}
        for poc_asym_id, pocket_mask in candidate_pockets.items():
            # Align pocket_i to true pocket
            rot, trans = get_optimal_transform(
                src_atoms=coord[pocket_mask].clone(),
                tgt_atoms=true_coord[true_pocket_mask],
                mask=true_coord_mask[true_pocket_mask],
            )
            # Transform predicted coordinates according to the aligment results
            aligned_pred_coord = apply_transform(coord.clone(), rot=rot, trans=trans)

            # Find the best ligand
            ordered_lig_asym_ids = [i for i in candidate_ligands]
            orderd_lig_masks = [candidate_ligands[i] for i in ordered_lig_asym_ids]
            aligned_lig_coords = torch.stack(
                [aligned_pred_coord[m] for m in orderd_lig_masks], dim=0
            )  # [N_lig, N_lig_atom, 3]

            if use_center_rmsd:
                mask = true_coord_mask[true_ligand_mask].bool()  # [N_lig_atom]
                aligned_lig_center = aligned_lig_coords[:, mask, :].mean(
                    dim=-2, keepdim=True
                )  # [N_lig, 1, 3]
                true_coord_center = true_coord[true_ligand_mask][mask, :].mean(
                    dim=-2, keepdim=True
                )  # [1, 3]
                per_lig_rmsd = rmsd(
                    aligned_lig_center,  # [N_lig, 1, 3]
                    expand_at_dim(
                        true_coord_center,
                        dim=0,
                        n=aligned_lig_coords.size(0),
                    ),
                    reduce=False,
                )  # [N_lig]
            else:
                per_lig_rmsd = rmsd(
                    aligned_lig_coords,
                    expand_at_dim(
                        true_coord[true_ligand_mask],
                        dim=0,
                        n=aligned_lig_coords.size(0),
                    ),
                    mask=true_coord_mask[true_ligand_mask],
                    reduce=False,
                )  # [N_lig]
            lig_rmsd, idx = per_lig_rmsd.min(dim=0)
            lig_asym_id = ordered_lig_asym_ids[idx]

            if lig_rmsd < best_results.get("rmsd", torch.inf):
                best_results = {
                    "rmsd": lig_rmsd,
                    "pocket_asym_id": poc_asym_id,
                    "ligand_asym_id": lig_asym_id,
                    "aligned_pred_coord": aligned_pred_coord,
                }
            if poc_asym_id == pocket_asym_id:
                # record the unpermuted result
                i = ordered_lig_asym_ids.index(ligand_asym_id)
                unpermuted_lig_rmsd = per_lig_rmsd[i].item()
                unpermuted_results = {
                    "rmsd": unpermuted_lig_rmsd,
                    "aligned_pred_coord": aligned_pred_coord,
                }

        # record stats
        per_sample_log_dict = {
            "is_permuted": best_results["pocket_asym_id"] != pocket_asym_id
            or best_results["ligand_asym_id"] != ligand_asym_id,
            "is_permuted_pocket": best_results["pocket_asym_id"] != pocket_asym_id,
            "is_permuted_ligand": best_results["ligand_asym_id"] != ligand_asym_id,
            "algo:no_permute": best_results["pocket_asym_id"] == pocket_asym_id
            and best_results["ligand_asym_id"] == ligand_asym_id,
        }
        improved_rmsd = (unpermuted_results["rmsd"] - best_results["rmsd"]).item()
        if improved_rmsd >= 1e-12:
            # better
            per_sample_log_dict.update(
                {
                    "algo:equivalent_permute": False,
                    "algo:worse_permute": False,
                    "algo:better_permute": True,
                    "algo:better_rmsd": improved_rmsd,
                }
            )
        elif improved_rmsd < 0:
            # worse
            per_sample_log_dict.update(
                {
                    "algo:equivalent_permute": False,
                    "algo:worse_permute": True,
                    "algo:better_permute": False,
                    "algo:worse_rmsd": -improved_rmsd,
                }
            )
        elif per_sample_log_dict["is_permuted"]:
            # equivalent
            per_sample_log_dict.update(
                {
                    "algo:equivalent_permute": True,
                    "algo:worse_permute": False,
                    "algo:better_permute": False,
                }
            )
        # atom indices to permute coordinates
        N_atom = aligned_pred_coord.size(-2)
        device = aligned_pred_coord.device
        atom_indices = torch.arange(N_atom, device=device)

        permute_asym_pair = [
            (best_results["pocket_asym_id"], pocket_asym_id),
            (best_results["ligand_asym_id"], ligand_asym_id),
        ]
        for asym_new, asym_old in permute_asym_pair:
            if asym_new == asym_old:
                continue
            # switch two chains
            ori_indices = atom_indices[atom_asym_id == asym_old]
            new_indices = atom_indices[atom_asym_id == asym_new]
            atom_indices[ori_indices.tolist()] = new_indices.clone()
            atom_indices[new_indices.tolist()] = ori_indices.clone()

        aligned_pred_coord = best_results.pop("aligned_pred_coord")[atom_indices, :]
        per_sample_log_dict["rmsd"] = best_results["rmsd"].item()

        return atom_indices, aligned_pred_coord, per_sample_log_dict

    N_sample = pred_coord.size(0)
    permute_pred_indices = []
    permuted_aligned_pred_coord = []
    sample_log_dicts = []
    for i in range(N_sample):
        atom_indices, aligned_pred_coord, per_sample_log_dict = (
            _find_protein_ligand_chains_for_one_sample(pred_coord[i])
        )
        permute_pred_indices.append(atom_indices)
        permuted_aligned_pred_coord.append(aligned_pred_coord)
        sample_log_dicts.append(per_sample_log_dict)

    permuted_aligned_pred_coord = torch.stack(permuted_aligned_pred_coord, dim=0)

    log_dict.update(
        traverse_and_aggregate(
            sample_log_dicts, aggregation_func=lambda x_list: sum(x_list) / N_sample
        )
    )

    # rmsd variance
    all_sample_rmsd = torch.tensor([x["rmsd"] for x in sample_log_dicts]).float()
    log_dict.update(
        {
            "rmsd_sample_std": all_sample_rmsd.std().item(),
            "rmsd_sample_gap": (all_sample_rmsd.max() - all_sample_rmsd.min()).item(),
        }
    )

    return permute_pred_indices, permuted_aligned_pred_coord, log_dict