# Copyright 2026 AlQuraishi Laboratory
#
# 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.

"""This module contains building blocks for template feature generation."""

import dataclasses
import logging
import traceback

import biotite.structure as struc
import numpy as np
import torch

from openfold3.core.data.primitives.featurization.structure import encode_one_hot
from openfold3.core.data.primitives.structure.template import TemplateSliceCollection
from openfold3.core.data.resources.residues import (
    STANDARD_RESIDUES_WITH_GAP_3,
    get_with_unknown_3_to_idx,
)
from openfold3.core.utils.all_atom_multimer import make_transform_from_reference
from openfold3.core.utils.geometry.vector import Vec3Array

logger = logging.getLogger(__name__)


@dataclasses.dataclass(frozen=False)
class OF3TemplateFeaturePrecursor:
    """Dataclass for storing information for OF3 template feature generation.

    Attributes:
        res_names (np.ndarray[str]):
            The names of the residues in the template structures. Has shape
            [n_templates, n_tokens].
        pseudo_beta_atom_coords (np.ndarray[float]):
            The coordinates of the pseudo beta atoms. Has shape
            [n_templates, n_tokens, 3].
        frame_atom_coords (np.ndarray[float]):
            The coordinates of the N, CA, and C atoms in this order along dim -2. Has
            shape [n_templates, n_tokens, 3, 3].
    """

    res_names: np.ndarray[str]
    pseudo_beta_atom_coords: np.ndarray[float]
    frame_atom_coords: np.ndarray[float]


def create_template_feature_precursor_of3(
    template_slice_collection: TemplateSliceCollection,
    n_templates: int,
    n_tokens: int,
) -> OF3TemplateFeaturePrecursor:
    """Generates set of precursor features for OF3 template feature generation.

    Args:
        template_slice_collection (TemplateSliceCollection):
            The collection of cropped template atom arrays per chain, per template.
        n_templates (int):
            Number of templates.
        n_tokens (int):
            Number of tokens in the target structure.

    Returns:
        OF3TemplateFeaturePrecursor:
            The precursor features for OF3 template feature generation. Includes
            residue names, pseudo beta atom coordinates, and N, CA, C atom coordinates.
    """
    res_names = np.full((n_templates, n_tokens), "GAP", dtype=np.dtype("U3"))
    pseudo_beta_atom_coords = np.full((n_templates, n_tokens, 3), np.nan, dtype=float)
    frame_atom_coords = np.full((n_templates, n_tokens, 3, 3), np.nan, dtype=float)

    # Iterate over chains then templates per chain
    for _, template_slices in template_slice_collection.template_slices.items():
        template_idx = 0
        for template_slice in template_slices:
            # TODO add support for these cases below Skip templates that fail to be
            # featurized these will be templates with composite residues like 4v2r
            # residue 195
            try:
                # Unpack template slice
                template_atom_array = template_slice.atom_array
                template_residue_repeats = template_slice.template_residue_repeats
                query_token_positions = template_slice.query_token_positions
                residue_starts = struc.get_residue_starts(template_atom_array)

                # Residue names and unmask corresponding tokens/template positions
                res_names[template_idx, query_token_positions] = np.repeat(
                    template_atom_array[residue_starts].res_name,
                    template_residue_repeats,
                )

                # Pseudo beta atom coordinates
                is_gly = template_atom_array.res_name == "GLY"
                is_ca = template_atom_array.atom_name == "CA"
                is_cb = template_atom_array.atom_name == "CB"
                is_pseudo_beta_atom = (is_gly & is_ca) | (~is_gly & is_cb)

                # Skip templates that have non-canonical residues with "non-canonical
                # C-beta atoms" or missing C-beta atoms
                # TODO: add handling for these
                # cases to be able to include them or skip them in a way that allows for
                # retaining the correct number of sampled templates
                if sum(is_pseudo_beta_atom) != len(residue_starts):
                    logger.debug(
                        "Skipping template with non-canonical/missing C-beta atoms."
                    )
                    continue

                pseudo_beta_atom_coords[template_idx, query_token_positions, :] = (
                    np.repeat(
                        template_atom_array[is_pseudo_beta_atom].coord,
                        template_residue_repeats,
                        axis=0,
                    )
                )

                # Frame (N, CA, C) atom coordinates
                is_n = template_atom_array.atom_name == "N"
                is_ca = template_atom_array.atom_name == "CA"
                is_c = template_atom_array.atom_name == "C"
                frame_atom_coords[template_idx, query_token_positions, 0, :] = (
                    np.repeat(
                        template_atom_array[is_n].coord,
                        template_residue_repeats,
                        axis=0,
                    )
                )
                frame_atom_coords[template_idx, query_token_positions, 1, :] = (
                    np.repeat(
                        template_atom_array[is_ca].coord,
                        template_residue_repeats,
                        axis=0,
                    )
                )
                frame_atom_coords[template_idx, query_token_positions, 2, :] = (
                    np.repeat(
                        template_atom_array[is_c].coord,
                        template_residue_repeats,
                        axis=0,
                    )
                )

                template_idx += 1
            except Exception as e:
                tb = traceback.format_exc()
                logger.warning(
                    "-" * 40
                    + "\n"
                    + f"Skipping template with exception: {str(e)}\n"
                    + f"Exception type: {type(e).__name__}\nTraceback: {tb}"
                    + "-" * 40
                )
                continue

    return OF3TemplateFeaturePrecursor(
        res_names=res_names,
        pseudo_beta_atom_coords=pseudo_beta_atom_coords,
        frame_atom_coords=frame_atom_coords,
    )


def create_template_restype(
    res_names: np.ndarray[str],
    template_pseudo_beta_mask: np.ndarray[float],
) -> torch.Tensor:
    """Creates the restype template feature for OF3.

    Args:
        res_names (np.ndarray[str]):
            The precursor features for OF3 template feature generation.
        template_pseudo_beta_mask (np.ndarray[float]):
            The mask for pseudo beta atoms. Has shape [n_templates, n_tokens].

    Returns:
        torch.Tensor:
            The restype template feature.
    """
    restype_index = torch.tensor(
        get_with_unknown_3_to_idx(res_names), dtype=torch.int64
    )
    template_restype = encode_one_hot(restype_index, len(STANDARD_RESIDUES_WITH_GAP_3))
    return template_restype.to(torch.int32)


def create_template_distogram(
    pseudo_beta_atom_coords: np.ndarray[float],
    pseudo_beta_mask: np.ndarray[float],
    multichain_pair_mask: torch.Tensor,
    min_bin: float = 3.25,
    max_bin: float = 50.75,
    n_bins: int = 39,
    inf_value: float = 1e8,
) -> torch.Tensor:
    """Creates the distogram template feature for OF3.

    Note: the pseudo_beta_mask is applied to the distogram to zero out masked
    tokens.

    Args:
        pseudo_beta_atom_coords (np.ndarray[float]):
            The coordinates of the pseudo beta atoms. Has shape
            [n_templates, n_tokens, 3].
        pseudo_beta_mask (np.ndarray[float]):
            Pseudo beta mask for tokens with an unresolved pseudo beta atom.
        multichain_pair_mask (torch.Tensor):
            Per-token-pair mask for inter-chain residue pair features.
        min_bin (float, optional):
            Bin lower bound. Defaults to 3.25.
        max_bin (float, optional):
            Bin upper bound. Defaults to 50.75.
        n_bins (int, optional):
            Number of bins between lower and upper bounds. Defaults to 39.
        inf_value (float, optional):
            Value to set last bin entries to. Defaults to 1e8.

    Returns:
        torch.Tensor:
            The distogram template feature [N_templates, N_token, N_token, N_bins].
    """
    distogram = np.sum(
        (
            pseudo_beta_atom_coords[..., None, :]
            - pseudo_beta_atom_coords[..., None, :, :]
        )
        ** 2,
        axis=-1,
        keepdims=True,
    )

    # Generate squared bin edges
    lower = np.linspace(min_bin, max_bin, n_bins) ** 2
    upper = np.concatenate(
        [lower[1:], np.array([inf_value], dtype=lower.dtype)], axis=-1
    )

    # Bin the distogram
    template_distogram = torch.tensor(
        ((distogram > lower) * (distogram < upper)).astype(distogram.dtype),
        dtype=torch.float,
    )

    return (
        template_distogram
        * (pseudo_beta_mask[..., None] * pseudo_beta_mask[..., None, :])[..., None]
        * multichain_pair_mask
    )


def create_template_unit_vector(
    frame_atom_coords: np.ndarray[float],
    backbone_frame_mask: np.ndarray[float],
    multichain_pair_mask: torch.Tensor,
) -> torch.Tensor:
    """Creates the unit vector template feature for OF3.

    Args:
        frame_atom_coords (np.ndarray[float]):
            The coordinates of the N, CA, and C atoms in this order along dim -2. Has
            shape [n_templates, n_tokens, 3, 3].
        backbone_frame_mask (np.ndarray[float]):
            The mask indicating for each token in each template if all backbone frame
            atoms are available. Has shape [n_templates, n_tokens].
        multichain_pair_mask (torch.Tensor):
            Per-token-pair mask for inter-chain residue pair features.

    Returns:
        torch.Tensor:
            The unit vector template feature [N_templates, N_token, N_token, 3].
    """

    # Convert nans to 0s and tensors to Vec3Arrays
    frame_atom_coords = torch.nan_to_num(
        torch.tensor(frame_atom_coords, dtype=torch.float32), nan=0.0
    )
    n_vec3array = Vec3Array.from_array(frame_atom_coords[:, :, 0, :])
    ca_vec3array = Vec3Array.from_array(frame_atom_coords[:, :, 1, :])
    c_vec3array = Vec3Array.from_array(frame_atom_coords[:, :, 2, :])

    # Create rigid frame and unit vector
    rigid = make_transform_from_reference(
        a_xyz=n_vec3array,
        b_xyz=ca_vec3array,
        c_xyz=c_vec3array,
    )
    unit_vector = (
        rigid[..., None]
        .inverse()
        .apply_to_point(rigid.translation[..., None, :])
        .normalized()
    )

    # Cast back to tensor and apply backbone frame mask
    masked_unit_vector = []
    for coord in unit_vector:
        # Create a 2D backbone frame mask
        backbone_frame_2d = (
            backbone_frame_mask[..., None] * backbone_frame_mask[..., None, :]
        )
        # Apply mask
        masked_coord = (coord * backbone_frame_2d)[..., None]
        masked_unit_vector.append(masked_coord)

    return torch.cat(masked_unit_vector, dim=-1) * multichain_pair_mask
