# Copyright (c) 2024 Chai Discovery, Inc.
# This source code is licensed under the Chai Discovery Community License
# Agreement (LICENSE.md) found in the root directory of this source tree.

import logging
from dataclasses import dataclass

import numpy as np
import torch
from einops import rearrange, repeat
from torch import Tensor

from chai_lab.data.features.feature_type import FeatureType
from chai_lab.data.features.generators.base import EncodingType, FeatureGenerator
from chai_lab.data.parsing.structure.entity_type import EntityType
from chai_lab.model.utils import get_asym_id_from_subchain_id
from chai_lab.utils.tensor_utils import cdist, tensorcode_to_string, und, und_self
from chai_lab.utils.typing import Bool, Float, Int, UInt8, typecheck

logger = logging.getLogger(__name__)


@typecheck
@dataclass
class ConstraintGroup:
    """
    Container for a token pair distance restraint (contact)
    """

    left_residue_subchain_id: str
    right_residue_subchain_id: str
    left_residue_index: int
    right_residue_index: int
    right_residue_name: str
    left_residue_name: str
    distance_threshold: float

    def get_left_and_right_asym_ids(
        self,
        token_subchain_id: UInt8[Tensor, "n 4"],
        token_asym_id: Int[Tensor, "n"],
    ):
        left_asym_id = get_asym_id_from_subchain_id(
            subchain_id=self.left_residue_subchain_id,
            source_pdb_chain_id=token_subchain_id,
            token_asym_id=token_asym_id,
        )
        right_asym_id = get_asym_id_from_subchain_id(
            subchain_id=self.right_residue_subchain_id,
            source_pdb_chain_id=token_subchain_id,
            token_asym_id=token_asym_id,
        )
        return left_asym_id, right_asym_id

    def __str__(self):
        return (
            f"ConstraintGroup(left_residue_subchain_id={self.left_residue_subchain_id}, "
            f"right_residue_subchain_id={self.right_residue_subchain_id}, "
            f"left_residue_index={self.left_residue_index}, "
            f"right_residue_index={self.right_residue_index}, "
            f"right_residue_name={self.right_residue_name}, "
            f"left_residue_name={self.left_residue_name}, "
            f"distance_threshold={self.distance_threshold})"
        )


class TokenDistanceRestraint(FeatureGenerator):
    def __init__(
        self,
        include_probability: float = 1.0,
        size: int | float = 0.33,
        min_dist: int | float = 10.0,
        max_dist: int | float = 30.0,
        coord_noise: float = 0.0,
        num_rbf_radii: int = 5,
        query_entity_types: list[EntityType] | None = None,
        key_entity_types: list[EntityType] | None = None,
    ):
        """Randomly sample inter-chain token distance restraints

        Parameters:
            include_probability: Probability with which to include restraints
                for a given example. i.e. if include probability is 0.75, then 25%
                of the time, we will not sample any restraints for an example.
            size:  Number of restraints to sample. If 0 < size < 1, then the number
                of restraints will be determined as geom(size), independently for each
                example.
            min_dist: Minimum distance to encode restraints for
            max_dist: Maximum distance to encode restraints for
            coord_noise: gaussian noise with mean 0 and variance coord_noise
                added to coordinates before sampling restraints.
            num_rbf_radii: Number of radii to use for the radial basis function
                embedding of restraints
            query_entity_types: Entity types to consider when sampling "query" tokens
                for restraints. Defaults to all entity types.
            key_entity_types: Entity types to consider when sampling "key" tokens
                for restraints. Defaults to all entity types.

        NOTE: We only sample restraints between tokens if one of the tokens is in
            the query entity types and the other is in the key entity types.
        """
        super().__init__(
            ty=FeatureType.TOKEN_PAIR,
            can_mask=False,
            encoding_ty=EncodingType.RBF,
            num_classes=num_rbf_radii,
            mult=1,
            ignore_index=-1.0,
        )
        self.ignore_idx = -1.0
        self.min_dist, self.max_dist = min_dist, max_dist
        self.coord_noise = coord_noise
        self.include_prob = include_probability
        self.size = size
        self.query_entity_types = torch.tensor(
            (
                [e.value for e in query_entity_types]
                if query_entity_types is not None
                else [e.value for e in EntityType]
            )
        ).long()
        self.key_entity_types = torch.tensor(
            [
                [e.value for e in key_entity_types]
                if key_entity_types is not None
                else [e.value for e in EntityType]
            ]
        ).long()

    def get_num_restraints(self, batch_size) -> list[int]:
        if 0 < self.size < 1:
            seles = np.random.geometric(self.size, size=batch_size)
            include_mask = np.random.uniform(size=batch_size) < self.include_prob
            seles[~include_mask] = 0
            return [int(x) for x in seles]
        return [int(self.size)] * batch_size

    def get_input_kwargs_from_batch(self, batch) -> dict:
        maybe_constraint_dicts = batch["inputs"].get("contact_constraints", [[None]])[0]
        contact_constraints = (
            [ConstraintGroup(**d) for d in maybe_constraint_dicts]
            if isinstance(maybe_constraint_dicts[0], dict)
            else None
        )
        return dict(
            atom_gt_coords=batch["inputs"]["atom_gt_coords"],
            atom_exists_mask=batch["inputs"]["atom_exists_mask"],
            token_asym_id=batch["inputs"]["token_asym_id"].long(),
            token_ref_atom_index=batch["inputs"]["token_ref_atom_index"].long(),
            token_exists_mask=batch["inputs"]["token_exists_mask"],
            token_entity_type=batch["inputs"]["token_entity_type"].long(),
            token_residue_index=batch["inputs"]["token_residue_index"].long(),
            token_residue_names=batch["inputs"]["token_residue_name"],
            token_subchain_id=batch["inputs"]["subchain_id"],
            constraints=contact_constraints,
        )

    def _sample_restraints(
        self,
        dists: Float[Tensor, "n n"],
        num_restraints: int,
    ):
        sampled_restraints = torch.full_like(dists, self.ignore_idx)
        # sample upper bound independently in range (min_dist, max_dist)
        # for each pair of tokens
        # We choose a random delta to upper bound all sampled distances with.
        # We do this because larger distance restraints are more likely to be
        # valid than smaller ones, and we try to reduce that bias here.
        delta = torch.rand(1) * (self.max_dist - self.min_dist)
        all_restraint_bounds = torch.rand_like(dists) * delta + self.min_dist
        all_valid_restraints = dists < all_restraint_bounds
        num_valid_restraints = int(all_valid_restraints.sum().item())
        if num_valid_restraints == 0 or num_restraints == 0:  # no restraints to add
            return sampled_restraints
        num_restraints = min(num_valid_restraints, num_restraints)
        # select random restraints and respective sampled bounds
        sampled_restraint_mask = all_restraint_bounds.new_zeros(
            num_valid_restraints, dtype=torch.bool
        )
        sampled_restraint_mask[:num_restraints] = True
        # select random restraints by shuffling
        sampled_restraint_mask = sampled_restraint_mask[
            torch.randperm(num_valid_restraints)
        ]

        # add the bounds/pairs that we sampled to the sampled restraint matrix
        flat_restraint_bounds = all_restraint_bounds[all_valid_restraints]
        flat_restraint_bounds[~sampled_restraint_mask] = self.ignore_idx
        sampled_restraints[all_valid_restraints] = flat_restraint_bounds

        return sampled_restraints

    @typecheck
    def _generate(
        self,
        atom_gt_coords: Float[Tensor, "b a 3"],
        atom_exists_mask: Bool[Tensor, "b a"],
        token_asym_id: Int[Tensor, "b n"],
        token_ref_atom_index: Int[Tensor, "b n"],
        token_exists_mask: Bool[Tensor, "b n"],
        token_entity_type: Int[Tensor, "b n"],
        token_residue_index: Int[Tensor, "b n"],
        token_residue_names: UInt8[Tensor, "b n 8"],
        token_subchain_id: UInt8[Tensor, "b n 4"],
        constraints: list[ConstraintGroup] | None = None,
    ) -> Tensor:
        try:
            if constraints is not None:
                assert atom_gt_coords.shape[0] == 1
                return self.generate_from_constraints(
                    token_asym_id=token_asym_id,
                    token_residue_index=token_residue_index,
                    token_residue_names=token_residue_names,
                    token_subchain_id=token_subchain_id,
                    constraints=constraints,
                )
        except Exception as e:
            logger.error(f"Error {e} generating distance constraints: {constraints}")

        return self._generate_from_batch(
            atom_gt_coords=atom_gt_coords,
            atom_exists_mask=atom_exists_mask,
            token_asym_id=token_asym_id,
            token_ref_atom_index=token_ref_atom_index,
            token_exists_mask=token_exists_mask,
            token_entity_type=token_entity_type,
        )

    @typecheck
    def _generate_from_batch(
        self,
        atom_gt_coords: Float[Tensor, "b a 3"],
        atom_exists_mask: Bool[Tensor, "b a"],
        token_asym_id: Int[Tensor, "b n"],
        token_ref_atom_index: Int[Tensor, "b n"],
        token_exists_mask: Bool[Tensor, "b n"],
        token_entity_type: Int[Tensor, "b n"],
    ) -> Tensor:
        batch_size = atom_gt_coords.shape[0]
        # create inter-chain contact mask
        valid_token_pair_mask = und_self(token_exists_mask, "b i, b j -> b i j")
        left_entity_type_mask = torch.any(
            (token_entity_type.unsqueeze(-1) - self.query_entity_types) == 0, dim=-1
        )
        right_entity_type_mask = torch.any(
            (token_entity_type.unsqueeze(-1) - self.key_entity_types) == 0, dim=-1
        )
        valid_entity_pair_mask = und(
            left_entity_type_mask, right_entity_type_mask, "b i, b j -> b i j"
        )
        diff_chain_mask = rearrange(token_asym_id, "b i -> b i 1") != rearrange(
            token_asym_id, "b j -> b 1 j"
        )
        ref_atom_mask = torch.gather(
            atom_exists_mask, dim=1, index=token_ref_atom_index
        )
        valid_token_ref_atom_mask = und_self(ref_atom_mask, "b i, b j -> b i j")
        valid_contact_mask = (
            valid_token_pair_mask
            & valid_entity_pair_mask
            & valid_token_ref_atom_mask
            & diff_chain_mask
        )

        # compute pairwise distances
        token_ref_atom_coords = torch.gather(
            atom_gt_coords, dim=1, index=repeat(token_ref_atom_index, "... -> ... 3")
        )
        # optionally add noise to coordinates before computing distances
        token_ref_atom_coords = (
            token_ref_atom_coords
            + torch.randn_like(token_ref_atom_coords) * self.coord_noise
        )
        inter_token_dists = cdist(token_ref_atom_coords)
        inter_token_dists[~valid_contact_mask] = self.max_dist + 1
        # compute contacts by (1) sampling an upper bound on the distance
        # and (2) selecting pairwise distances below the threshold
        num_to_include = self.get_num_restraints(batch_size)
        restraint_mats = [
            self._sample_restraints(inter_token_dists[i], n)
            for i, n in enumerate(num_to_include)
        ]
        encoded_feat = torch.stack(restraint_mats, dim=0)
        return self.make_feature(encoded_feat.unsqueeze(-1))

    @typecheck
    def generate_from_constraints(
        self,
        token_asym_id: Int[Tensor, "1 n"],
        token_residue_index: Int[Tensor, "1 n"],
        token_residue_names: UInt8[Tensor, "1 n 8"],
        token_subchain_id: UInt8[Tensor, "1 n 4"],
        constraints: list[ConstraintGroup],
    ) -> Tensor:
        logger.info(f"Generating distance feature from constraints: {constraints}")
        n, device = token_asym_id.shape[1], token_asym_id.device
        constraint_mat = torch.full(
            (n, n), fill_value=self.ignore_idx, device=device, dtype=torch.float32
        )
        for constraint_group in constraints:
            left_residue_asym_id, right_residue_asym_id = (
                constraint_group.get_left_and_right_asym_ids(
                    token_subchain_id=rearrange(token_subchain_id, "1 ... -> ..."),
                    token_asym_id=rearrange(token_asym_id, "1 ... -> ..."),
                )
            )
            constraint_mat = self.add_distance_constraint(
                constraint_mat=constraint_mat,
                token_asym_id=rearrange(token_asym_id, "1 ... -> ..."),
                token_residue_index=rearrange(token_residue_index, "1 ... -> ..."),
                token_residue_names=rearrange(token_residue_names, "1 ... -> ..."),
                left_residue_asym_id=left_residue_asym_id,
                right_residue_asym_id=right_residue_asym_id,
                left_residue_index=constraint_group.left_residue_index,
                right_residue_index=constraint_group.right_residue_index,
                right_residue_name=constraint_group.right_residue_name,
                left_residue_name=constraint_group.left_residue_name,
                distance_threshold=constraint_group.distance_threshold,
            )
        # encode and apply mask
        constraint_mat = repeat(constraint_mat, "i j -> 1 i j 1")
        return self.make_feature(constraint_mat)

    @typecheck
    def add_distance_constraint(
        self,
        constraint_mat: Float[Tensor, "n n"],
        token_asym_id: Int[Tensor, "n"],
        token_residue_index: Int[Tensor, "n"],
        token_residue_names: UInt8[Tensor, "n 8"],
        # asym id of the chain that binds in the pocket
        left_residue_asym_id: int,
        right_residue_asym_id: int,
        left_residue_index: int,
        right_residue_index: int,
        right_residue_name: str,
        left_residue_name: str,
        distance_threshold: float,
    ):
        left_asym_mask = token_asym_id == left_residue_asym_id
        right_asym_mask = token_asym_id == right_residue_asym_id
        left_index_mask = token_residue_index == left_residue_index
        right_index_mask = token_residue_index == right_residue_index
        left_residue_mask = left_asym_mask & left_index_mask
        right_residue_mask = right_asym_mask & right_index_mask
        # restraint should point to single residue pair
        assert torch.sum(left_residue_mask) == 1, (
            f"Expected unique residue but found {torch.sum(left_residue_mask)}\n"
            f"{left_residue_asym_id=}, {left_residue_index=}, {left_residue_name=}"
        )
        assert torch.sum(right_residue_mask) == 1, (
            f"Expected unique residue but found {torch.sum(right_residue_mask)}\n"
            f"{right_residue_asym_id=}, {right_residue_index=}, {right_residue_name=}"
        )
        # make sure the residue names in the constraint match the
        # ones we parsed
        left_res_name = token_residue_names[left_residue_mask]
        right_res_name = token_residue_names[right_residue_mask]
        expected_res_name = tensorcode_to_string(rearrange(left_res_name, "1 l -> l"))
        assert expected_res_name == left_residue_name, (
            f"Expected residue name {expected_res_name} but got " f"{left_residue_name}"
        )
        expected_res_name = tensorcode_to_string(rearrange(right_res_name, "1 l -> l"))
        assert expected_res_name == right_residue_name, (
            f"Expected residue name {expected_res_name} but got "
            f"{right_residue_name}"
        )
        # add constraint
        # NOTE: feature is *not* symmetric
        constraint_mat[left_residue_mask, right_residue_mask] = distance_threshold
        return constraint_mat
