import logging
from dataclasses import dataclass
from itertools import chain

import torch
from einops import repeat
from torch import Tensor

from chai_lab.data.dataset.structure import utils
from chai_lab.data.dataset.structure.all_atom_structure_context import (
    AllAtomStructureContext,
)
from chai_lab.data.dataset.structure.utils import (
    backbone_atoms_all_present,
    backbone_atoms_indices,
    get_centre_atom_index,
    get_reference_atom_index,
)
from chai_lab.data.parsing.structure.all_atom_entity_data import AllAtomEntityData
from chai_lab.data.parsing.structure.entity_type import EntityType
from chai_lab.data.parsing.structure.residue import ConformerData, Residue
from chai_lab.data.residue_constants import standard_residue_pdb_codes
from chai_lab.data.sources.rdkit import (
    RefConformerGenerator,
    conformer_data_to_rdkit_mol,
)
from chai_lab.utils.tensor_utils import string_to_tensorcode, unique_indexes
from chai_lab.utils.typing import Bool, Float, Int, typecheck

logger = logging.getLogger(__name__)


# jaxtyping on residue-level objects is extremely slow.
@dataclass(frozen=True)
class TokenSpan:
    restype: Int[Tensor, "n_tokens"]
    residue_index: Int[Tensor, "n_tokens"]
    centre_atom_index: Int[Tensor, "n_tokens"]
    reference_atom_index: Int[Tensor, "n_tokens"]
    backbone_frame_mask: Bool[Tensor, "n_tokens"]
    backbone_frame_index: Int[Tensor, "n_tokens 3"]
    atom_gt_coords: Float[Tensor, "n_atoms 3"]
    atom_exists_mask: Bool[Tensor, "n_atoms"]
    atom_token_index: Int[Tensor, "n_atoms"]
    ref_pos: Float[Tensor, "n_atoms 3"]
    ref_mask: Bool[Tensor, "n_atoms"]
    ref_element: Int[Tensor, "n_atoms"]
    ref_charge: Int[Tensor, "n_atoms"]
    atom_names: list[str]
    # Consistent atom ordering witin each token
    atom_within_token_indices: Int[Tensor, "n_atoms"]
    residue_names: list[str]
    symmetries: Int[Tensor, "n_atoms n_symm"]
    b_factor_or_plddt: Float[Tensor, "n_tokens"]

    @classmethod
    def concatenate(cls, spans: list["TokenSpan"]) -> "TokenSpan":
        # offset bond indices:
        tokens_per_span = torch.tensor([span.restype.shape[0] for span in spans])
        token_count = torch.cumsum(tokens_per_span, dim=0).roll(1, 0)
        token_count[0] = 0

        # offsets indices of centre atoms:
        atoms_per_span = torch.tensor(
            [span.atom_exists_mask.shape[0] for span in spans]
        )
        atom_offsets = torch.cumsum(atoms_per_span, dim=0).roll(1, 0)
        atom_offsets[0] = 0

        centre_atom_index = torch.cat(
            [
                span.centre_atom_index + offset
                for span, offset in zip(spans, atom_offsets)
            ]
        )
        reference_atom_index = torch.cat(
            [
                span.reference_atom_index + offset
                for span, offset in zip(spans, atom_offsets)
            ]
        )

        atom_token_index = (
            torch.cumsum(
                torch.cat([x.atom_token_index for x in spans]),
                dim=0,
                dtype=torch.int,
            )
            - 1
        )
        backbone_frame_index = torch.cat(
            [
                span.backbone_frame_index + offset
                for span, offset in zip(spans, atom_offsets)
            ]
        )

        # concatenate symmetric permutations at the atom level
        # make sure that trailing shape is the same
        # NOTE: we store the *local* permutation indices, not the global ones
        # i.e. the permutation indices are relative to the residue
        atom_symms = [span.symmetries for span in spans]
        max_symms = max(x.shape[-1] for x in atom_symms)
        atom_symms = [
            torch.nn.functional.pad(x, (0, max_symms - x.shape[-1]), value=-1)
            for x in atom_symms
        ]
        return cls(
            restype=torch.cat([x.restype for x in spans]),
            residue_index=torch.cat([x.residue_index for x in spans]),
            centre_atom_index=centre_atom_index,
            reference_atom_index=reference_atom_index,
            backbone_frame_mask=torch.cat([x.backbone_frame_mask for x in spans]),
            backbone_frame_index=backbone_frame_index,
            atom_gt_coords=torch.cat([x.atom_gt_coords for x in spans]),
            atom_exists_mask=torch.cat([x.atom_exists_mask for x in spans]),
            atom_token_index=atom_token_index,
            ref_pos=torch.cat([x.ref_pos for x in spans]),
            ref_mask=torch.cat([x.ref_mask for x in spans]),
            ref_element=torch.cat([x.ref_element for x in spans]),
            ref_charge=torch.cat([x.ref_charge for x in spans]),
            atom_names=list(chain.from_iterable([x.atom_names for x in spans])),
            atom_within_token_indices=torch.cat(
                [x.atom_within_token_indices for x in spans]
            ),
            residue_names=list(chain.from_iterable([x.residue_names for x in spans])),
            symmetries=torch.cat(atom_symms, dim=0),
            b_factor_or_plddt=torch.cat([x.b_factor_or_plddt for x in spans]),
        )


class AllAtomResidueTokenizer:
    ref_conformer_generator: RefConformerGenerator

    def __init__(self, ref_conformer_generator: RefConformerGenerator):
        self.ref_conformer_generator = ref_conformer_generator

    def tokenize_residue(
        self,
        residue: Residue,
        entity_type: EntityType,
    ) -> TokenSpan | None:
        ref_conformer_data = self._get_ref_conformer_data(residue)
        if ref_conformer_data.num_atoms == 0:
            # avoid dealing with empty tensors in downstream processing
            # this should only happen when residue is sole hydrogen
            # or when residue code is not in CCD dictionary and
            # the residue has 0 coords in the PDB structure
            logger.warning(
                f"skipping residue {residue.name} {residue.label_seq} as reference conformer has 0 heavy atoms"
            )
            return None

        # Keep only the atoms from the ground truth conformer that are present in
        # reference conformer.
        #
        # If we don't have a reference conformer, we fall back to using the ground truth
        # conformer names, i.e. we keep all atoms in the ground truth conformer.
        # When a true conformer data is not provided, use reference conformer directly
        gt_conformer_data = residue.conformer_data

        if gt_conformer_data is not None:
            atom_gt_coords, atom_exists_mask = gt_conformer_data.gather_atom_positions(
                ref_conformer_data.atom_names
            )
        else:
            atom_gt_coords = ref_conformer_data.position
            atom_exists_mask = torch.ones(
                atom_gt_coords.shape[0], dtype=torch.bool, device=atom_gt_coords.device
            )

        # Tokenization is by residue if it is a standard amino acid or standard
        # nucleotide; all ligands and all modified residues are tokenized per atom.
        tokenize_fn = (
            self._tokenize_per_residue
            if (
                residue.name in standard_residue_pdb_codes
                and entity_type != EntityType.LIGAND
            )
            else self._tokenize_per_atom
        )

        return tokenize_fn(
            restype=torch.tensor([residue.restype], dtype=torch.int),
            residue_index=torch.tensor([residue.residue_index], dtype=torch.int),
            atom_gt_coords=atom_gt_coords,
            atom_exists_mask=atom_exists_mask,
            ref_pos=ref_conformer_data.position,
            ref_mask=torch.ones_like(atom_exists_mask),
            ref_element=ref_conformer_data.element,
            ref_charge=ref_conformer_data.charge,
            atom_names=ref_conformer_data.atom_names,
            residue_name=residue.name,
            bonds=ref_conformer_data.bonds,
            symmetries=ref_conformer_data.symmetries,
            b_factor_or_plddt=torch.tensor([residue.b_factor_or_plddt]),
        )

    @staticmethod
    def filter_atom_symmetries(
        symmetries: Int[Tensor, "n_atoms n_symm"],
        atom_exists_mask: Bool[Tensor, "n_atoms"],
    ) -> Int[Tensor, "n_atoms filtered_n_symm"]:
        n_atoms, _ = symmetries.shape

        # Create a mask for non-trivial symmetries
        atom_indices = torch.arange(n_atoms).unsqueeze(-1)
        non_trivial_symmetries = (symmetries >= 0) & (symmetries != atom_indices)

        masked_atoms = ~atom_exists_mask.unsqueeze(-1)

        # Check if any of the masked-out atoms have non-trivial symmetries
        violations = torch.any(masked_atoms & non_trivial_symmetries, dim=1)

        # If any invalid symmetries are found, replace with identity permutation
        if torch.any(violations):
            return atom_indices

        # Otherwise, return the original symmetries
        return symmetries

    # jaxtyping on residue-level objects is very slow,
    # use for debug only
    # @typecheck
    def _tokenize_per_residue(
        self,
        restype: Int[Tensor, "n_tokens"],
        residue_index: Int[Tensor, "n_tokens"],
        atom_gt_coords: Float[Tensor, "n_atoms 3"],
        atom_exists_mask: Bool[Tensor, "n_atoms"],
        ref_pos: Float[Tensor, "n_atoms 3"],
        ref_mask: Bool[Tensor, "n_atoms"],
        ref_element: Int[Tensor, "n_atoms"],
        ref_charge: Int[Tensor, "n_atoms"],
        atom_names: list[str],
        residue_name: str,
        bonds: list[tuple[int, int]],
        symmetries: Int[Tensor, "n_atoms n_symm"],
        b_factor_or_plddt: Float[Tensor, "n_tokens"],
    ) -> TokenSpan:
        centre_atom_index = get_centre_atom_index(
            atom_names,
            residue_name,
        )
        reference_atom_index = get_reference_atom_index(
            atom_names,
            residue_name,
        )
        backbone_frame_mask = backbone_atoms_all_present(
            atom_names,
            residue_name,
        )
        backbone_indices = backbone_atoms_indices(atom_names, residue_name).unsqueeze(0)

        # to 1 token
        atom_token_index = torch.zeros_like(atom_exists_mask, dtype=torch.int)
        atom_token_index[0] = 1

        residue_names = [residue_name]

        # Find atom ordering; these should always be available because per residue
        # tokenization works only on standard residues.
        atom_within_token_index = atom_names_to_atom37_indices(
            atom_names=atom_names,
            residue_name=residue_name,
        )

        return TokenSpan(
            restype=restype,
            residue_index=residue_index,
            centre_atom_index=centre_atom_index,
            reference_atom_index=reference_atom_index,
            backbone_frame_mask=backbone_frame_mask,
            backbone_frame_index=backbone_indices,
            atom_gt_coords=atom_gt_coords,
            atom_exists_mask=atom_exists_mask,
            atom_token_index=atom_token_index,
            ref_pos=ref_pos,
            ref_mask=ref_mask,
            ref_element=ref_element,
            ref_charge=ref_charge,
            atom_names=atom_names,
            atom_within_token_indices=atom_within_token_index,
            residue_names=residue_names,
            symmetries=self.filter_atom_symmetries(symmetries, atom_exists_mask),
            b_factor_or_plddt=b_factor_or_plddt,
        )

    # jaxtyping on residue-level objects is very slow,
    # use for debug only
    # @typecheck
    def _tokenize_per_atom(
        self,
        restype: Int[Tensor, "n_tokens"],
        residue_index: Int[Tensor, "n_tokens"],
        atom_gt_coords: Float[Tensor, "n_atoms 3"],
        atom_exists_mask: Bool[Tensor, "n_atoms"],
        ref_pos: Float[Tensor, "n_atoms 3"],
        ref_mask: Bool[Tensor, "n_atoms"],
        ref_element: Int[Tensor, "n_atoms"],
        ref_charge: Int[Tensor, "n_atoms"],
        atom_names: list[str],
        residue_name: str,
        bonds: list[tuple[int, int]],
        symmetries: Int[Tensor, "n_atoms n_symm"],
        b_factor_or_plddt: Float[Tensor, "n_tokens"],
    ) -> TokenSpan:
        # to n_atoms tokens
        n_atoms = atom_gt_coords.shape[0]
        restype = repeat(restype, "1 -> a", a=n_atoms)
        residue_index = repeat(residue_index, "1 -> a", a=n_atoms)
        b_factor_or_plddt = repeat(b_factor_or_plddt, "1 -> a", a=n_atoms)

        # centre of the token is the first and only atom in each token
        # when tokenizing per-atom
        centre_atom_index = torch.arange(n_atoms, dtype=torch.int)
        reference_atom_index = torch.arange(n_atoms, dtype=torch.int)
        backbone_frame_mask = torch.zeros((n_atoms,), dtype=torch.bool)
        backbone_indices = (
            torch.arange(n_atoms, dtype=torch.int).unsqueeze(1).expand(-1, 3)
        )

        atom_token_index = torch.ones_like(atom_exists_mask, dtype=torch.int)

        residue_names = [residue_name] * n_atoms

        # Each atom is alone in its own token
        atom_within_token_index = torch.zeros(n_atoms, dtype=torch.int)

        return TokenSpan(
            restype=restype,
            residue_index=residue_index,
            centre_atom_index=centre_atom_index,
            reference_atom_index=reference_atom_index,
            backbone_frame_mask=backbone_frame_mask,
            backbone_frame_index=backbone_indices,
            atom_gt_coords=atom_gt_coords,
            atom_exists_mask=atom_exists_mask,
            atom_token_index=atom_token_index,
            ref_pos=ref_pos,
            ref_mask=ref_mask,
            ref_element=ref_element,
            ref_charge=ref_charge,
            atom_names=atom_names,
            atom_within_token_indices=atom_within_token_index,
            residue_names=residue_names,
            symmetries=self.filter_atom_symmetries(symmetries, atom_exists_mask),
            b_factor_or_plddt=b_factor_or_plddt,
        )

    def tokenize_entity(
        self, entity_data: AllAtomEntityData
    ) -> AllAtomStructureContext | None:
        return self.tokenize_entities([entity_data])[0]

    def tokenize_entities(
        self,
        entities_data: list[AllAtomEntityData],
    ) -> list[AllAtomStructureContext | None]:
        sym_ids = _make_sym_ids([x.entity_id for x in entities_data])

        return [
            self._tokenize_entity(
                entity_data,
                chain_id=idx + 1,
                sym_id=sym_id,
            )
            for idx, (entity_data, sym_id) in enumerate(zip(entities_data, sym_ids))
        ]

    def _tokenize_entity(
        self,
        entity_data: AllAtomEntityData,
        chain_id: int = 1,
        sym_id: int = 1,
    ) -> AllAtomStructureContext | None:
        tokenized_residues = [
            self.tokenize_residue(residue, entity_data.entity_type)
            for residue in entity_data.residues
        ]

        valid_residues = [x for x in tokenized_residues if x is not None]
        if len(valid_residues) == 0:
            return None

        tokens = TokenSpan.concatenate(valid_residues)

        num_tokens = tokens.restype.shape[0]
        token_index = torch.arange(num_tokens, dtype=torch.int)

        # mask indicating if a token has >=1 atom with known coordinates
        token_exists_mask = (tokens.atom_token_index == token_index[..., None]).sum(
            dim=-1
        ) > 0

        # checks on atom mask and positions:
        # max 1 atom per-example has zero coordinates
        if (
            torch.sum(
                torch.all(tokens.atom_gt_coords[tokens.atom_exists_mask] == 0, dim=-1)
            )
            > 1
        ):
            raise ValueError(
                f"Zero coordinates found in unmasked atoms for {entity_data.pdb_id}"
            )

        # construct asym_id, entity_id, sym_id
        asym_id = chain_id
        entity_id = entity_data.entity_id

        # Create unique ids to identify atoms which belong to same residue in same chain
        # here assume we featurize a single chain
        atom_residue_index = torch.gather(
            tokens.residue_index,
            dim=0,
            index=tokens.atom_token_index.long(),
        )

        atom_ref_space_uid = atom_residue_index

        residue_names = tokens.residue_names

        match entity_data.entity_type:
            case EntityType.PROTEIN:
                if tokens.residue_index[0].item() != 0:
                    logger.error(
                        f"Protein residue index should start at zero, {entity_data}"
                    )

                if not torch.all(torch.diff(tokens.residue_index) <= 1):
                    logger.error(
                        f"Protein residue index should be contiguous (no gaps), {entity_data}"
                    )

                _, unique_indices = unique_indexes(tokens.residue_index)
                res_seq = [residue_names[i.item()] for i in unique_indices]
                if res_seq != entity_data.full_sequence:
                    logger.error(
                        f"Protein residue names should match entity data full sequence, {entity_data}"
                    )

        return AllAtomStructureContext(
            # token-level
            token_residue_type=tokens.restype,
            token_residue_index=tokens.residue_index,
            token_centre_atom_index=tokens.centre_atom_index,
            token_ref_atom_index=tokens.reference_atom_index,
            token_index=token_index,
            token_exists_mask=token_exists_mask,
            token_backbone_frame_mask=tokens.backbone_frame_mask,
            token_backbone_frame_index=tokens.backbone_frame_index,
            token_asym_id=_id_to_token_tensor(asym_id, num_tokens),
            token_entity_id=_id_to_token_tensor(entity_id, num_tokens),
            token_sym_id=_id_to_token_tensor(sym_id, num_tokens),
            token_entity_type=entity_type_to_tensor(
                entity_data.entity_type,
                num_tokens,
            ),
            # token res name is padded to 8 characters
            token_residue_name=torch.stack(
                [string_to_tensorcode(x, 8) for x in residue_names],
                dim=0,
            ),
            token_b_factor_or_plddt=tokens.b_factor_or_plddt,
            # atom-level
            atom_token_index=tokens.atom_token_index,
            atom_within_token_index=tokens.atom_within_token_indices,
            atom_ref_pos=tokens.ref_pos,
            atom_ref_mask=tokens.ref_mask,
            atom_ref_element=tokens.ref_element,
            atom_ref_charge=tokens.ref_charge,
            atom_ref_name=tokens.atom_names,
            atom_ref_name_chars=_atom_names_to_tensor(tokens.atom_names),
            atom_ref_space_uid=atom_ref_space_uid,
            atom_is_not_padding_mask=torch.ones_like(
                tokens.atom_exists_mask,
                dtype=torch.bool,
            ),
            # supervision only
            atom_gt_coords=tokens.atom_gt_coords,
            atom_exists_mask=tokens.atom_exists_mask,
            # structure-only
            pdb_id=repeat(
                # PDB ids are only 4 characters long, but AFDB ids can be longer
                string_to_tensorcode(entity_data.pdb_id, pad_to_length=32),
                "length -> num_tokens length",
                num_tokens=num_tokens,
            ),
            source_pdb_chain_id=repeat(
                string_to_tensorcode(entity_data.source_pdb_chain_id, pad_to_length=4),
                "length -> num_tokens length",
                num_tokens=num_tokens,
            ),
            subchain_id=repeat(
                string_to_tensorcode(entity_data.subchain_id, pad_to_length=4),
                "length -> num_tokens length",
                num_tokens=num_tokens,
            ),
            resolution=torch.tensor(
                [entity_data.resolution],
                dtype=torch.float32,
            ),
            is_distillation=torch.tensor(
                [entity_data.is_distillation],
                dtype=torch.bool,
            ),
            symmetries=tokens.symmetries,
        )

    def _get_ref_conformer_data(self, residue: Residue) -> ConformerData:
        """
        Returns the reference conformer data for the residue. We determine the reference
        conformer according to the following logic:
        1. conformer_generator is available and a reference
           conformer exists for the residue name => we return the cached reference
           conformer via the conformer generator
        2. conformer_generator is available and a smiles is given for the residue =>
            we generate a reference conformer using Rdkit via the conformer generator
        3. conformer_generator is available, the reference conformer can't be
            found and no smiles is given => we convert the Residue to an RDKit molecule
            and load full conformer data with the residue's atom positions as coordinates.
        4. conformer generator is not available => we set reference conformer to
            the ground truth conformer
        """
        # The reference conformer tells us:
        # - which atoms we should expect in this ligand / residue, and how many of them
        # - what are the ideal coordinates of these atoms if the ligand or residue was
        #   assembled alone in the void
        ref_conformer = self.ref_conformer_generator.get(residue.name)

        if ref_conformer is not None:
            if residue.name in standard_residue_pdb_codes:
                return ref_conformer
            else:
                return ref_conformer.center_random_augment()

        # When we can't find a reference conformer, and a smiles is given,
        # generate a reference conformer using rdkit
        if residue.smiles is not None:
            logger.info(
                f"Generating ref conformer for {residue.name}, {residue.smiles}"
            )
            return self.ref_conformer_generator.generate(residue.smiles)

        # When we can't find a reference conformer, attempt to use the ground
        # truth conformer data as the reference conformer.
        logger.warning(
            f"No reference conformer found for residue {residue.name},"
            "using training example conformer"
        )
        assert residue.conformer_data is not None

        try:
            # Rather than just setting the reference conformer to the ground truth, we
            # make a fake RDKit molecule from the ground truth data and then convert
            # back into a conformer data so that we can extract inter-atom aymmetries
            # bond and info
            rdkit_mol = conformer_data_to_rdkit_mol(residue.conformer_data)
            gt_conformer = RefConformerGenerator._load_ref_conformer_from_rdkit(
                rdkit_mol
            )
        except Exception as e:
            # Occasionally _load_ref_conformer_from_rdkit fails on unknown ligands e.g.
            # rdkit.Chem.rdchem.AtomValenceException's can be raised or ValueError:
            # can't infer bonds for Ligand. due to inexact connectivity.
            logger.warning(
                f"Caught error for {residue.name=} while loading reference conformer "
                f"from RDKit, {(type(e).__name__)}. Using ground truth conformer instead."
            )
            gt_conformer = residue.conformer_data

        return gt_conformer.center_random_augment()


@typecheck
def _atom_names_to_tensor(atom_names: list[str]) -> Int[Tensor, "n_atoms 4"]:
    ords = torch.tensor(
        [[ord(c) - 32 for c in atom_name.ljust(4, " ")] for atom_name in atom_names],
        dtype=torch.int,
    )
    return ords[:, :4]


@typecheck
def _id_to_token_tensor(id: int, num_tokens: int) -> Int[Tensor, "n"]:
    return id * torch.ones((num_tokens,), dtype=torch.int)


@typecheck
def entity_type_to_tensor(entity_type: EntityType, num_tokens: int) -> Int[Tensor, "n"]:
    return torch.full((num_tokens,), fill_value=entity_type.value, dtype=torch.int)


def _make_sym_ids(entity_ids_per_chain: list[int]) -> list[int]:
    entities_dict: dict[int, int] = dict()
    sym_ids = []

    for entity_id in entity_ids_per_chain:
        sym_id = entities_dict.get(entity_id, 0)
        sym_ids.append(sym_id)
        entities_dict[entity_id] = sym_id + 1

    return sym_ids


def atom_names_to_atom37_indices(
    atom_names: list[str], residue_name: str
) -> Int[Tensor, "n_atoms"]:
    """
    Returns a tensor of indices into the token-level atom names.
    """
    # Proteins use the atom37 ordering and indexing
    # nucleotides use the 36 atom ordering and indexing
    # - DNA is written as DA DG DC DT
    # - RNA is given as A G C U

    precomputed_idces = utils.atom_37_atom_indices()

    if residue_name == "UNK":
        retval = torch.arange(len(atom_names), dtype=torch.int)

    elif residue_name in standard_residue_pdb_codes:
        idx = [precomputed_idces[(residue_name, atom_name)] for atom_name in atom_names]
        retval = torch.tensor(idx, dtype=torch.int)
    else:
        raise ValueError(
            f"Unknown residue name {residue_name} (atom names: {atom_names})"
        )

    assert retval.max() <= 36, f"Out of bounds ordering {atom_names} in {residue_name}"
    return retval
