# Copyright 2025 Dhruv Nair. All rights reserved.
#
# 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.

from typing import List, Tuple, Union

import torch

from diffusers.utils import logging
from diffusers.modular_pipelines import ModularPipeline, ModularPipelineBlocks, PipelineState
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam


logger = logging.get_logger(__name__)

# Feature widths the checkpoint's token initializer was trained with, from
# rfd3/configs/model/components/rfd3_net.yaml. They set the input width of the embedding layers,
# so a mismatch raises on layer shape rather than silently degrading the design.
_TOKEN_1D_FEATURES = {"ref_motif_token_type": 3, "restype": 32, "ref_plddt": 1, "is_non_loopy": 1}
_ATOM_1D_FEATURES = {
    "ref_atom_name_chars": 256,
    "ref_element": 128,
    "ref_charge": 1,
    "ref_mask": 1,
    "ref_is_motif_atom_with_fixed_coord": 1,
    "ref_is_motif_atom_unindexed": 1,
    "has_zero_occupancy": 1,
    "ref_pos": 3,
    "ref_atomwise_rasa": 3,
    "active_donor": 1,
    "active_acceptor": 1,
    "is_atom_level_hotspot": 1,
}

# Every token is padded to this many atom slots, so the coordinate tensors the model sees are
# atom level with L = _N_ATOMS_PER_TOKEN * n_residues.
_N_ATOMS_PER_TOKEN = 14


def build_design_features(length: int, diffusion_batch_size: int, sigma_data: float) -> dict:
    """
    Build the foundry feature dict for an unconditional design of `length` residues.

    The token initializer embeds real chemical and positional features; there is no meaningful
    zero substitute for them, and no API in foundry that turns a length into features in one
    call. This mirrors what `rfd3.engine.RFD3InferenceEngine` does, minus the hydra app and the
    checkpoint, by driving the same specification and transform pipeline directly.

    Returns:
        The transformed example, carrying `feats` and `coord_atom_lvl_to_be_noised`.
    """
    from rfd3.inference.input_parsing import DesignInputSpecification
    from rfd3.transforms.pipelines import build_atom14_base_pipeline

    # `length` and `contig` are mutually exclusive in the specification: passing `contig` without
    # a structure input fails validation, so unconditional designs go through `length`.
    spec = DesignInputSpecification(length=str(length))
    data = spec.to_pipeline_input(example_id=f"rfd3_{length}")

    pipeline = build_atom14_base_pipeline(
        is_inference=True,
        diffusion_batch_size=diffusion_batch_size,
        sigma_data=sigma_data,
        central_atom="CB",
        n_atoms_per_token=_N_ATOMS_PER_TOKEN,
        generate_conformers=True,
        provide_reference_conformer_when_unmasked=True,
        ground_truth_conformer_policy="IGNORE",
        use_element_for_atom_names_of_atomized_tokens=True,
        token_1d_features=_TOKEN_1D_FEATURES,
        atom_1d_features=_ATOM_1D_FEATURES,
    )
    return pipeline(data)


def parse_contig_string(contig_str: str) -> Tuple[int, List[Tuple[int, int]]]:
    """
    Parse contig specification string.

    Supports formats like:
    - "100" -> 100 residues to design
    - "50-100" -> random length between 50-100
    - "A10-25/50" -> motif from chain A residues 10-25, plus 50 designed

    Returns:
        total_length: Total protein length
        motif_ranges: List of (start, end) for motif residues (0-indexed)
    """
    parts = contig_str.split("/")
    total_length = 0
    motif_ranges = []

    for part in parts:
        part = part.strip()
        if not part:
            continue

        if part[0].isalpha():
            chain = part[0]
            residue_spec = part[1:]
            if "-" in residue_spec:
                start, end = map(int, residue_spec.split("-"))
            else:
                start = end = int(residue_spec)
            motif_len = end - start + 1
            motif_ranges.append((total_length, total_length + motif_len))
            total_length += motif_len
        else:
            if "-" in part:
                min_len, max_len = map(int, part.split("-"))
                add_len = (min_len + max_len) // 2
            else:
                add_len = int(part)
            total_length += add_len

    return total_length, motif_ranges


class RFDiffusionInputStep(ModularPipelineBlocks):
    """
    Input processing step for RFDiffusion.

    Parses contigs to prepare features for structure generation.
    """

    model_name = "rfdiffusion"

    @property
    def description(self) -> str:
        return (
            "Input processing step that:\n"
            "  1. Parses contig specification to determine protein length and design regions\n"
            "  2. Generates masks for motif positions\n"
        )

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam(
                "contigs",
                required=True,
                type_hint=Union[str, List[str]],
                description="Contig specification defining design regions (e.g., '100' or 'A10-25/50-100')",
            ),
            InputParam(
                "input_xyz",
                type_hint=torch.Tensor,
                description="Input coordinates for motif residues [N_motif, 3]",
            ),
        ]

    @property
    def expected_components(self) -> List[ComponentSpec]:
        return [
            ComponentSpec("scheduler", description="RFDiffusion3 EDM scheduler"),
        ]

    @property
    def intermediate_outputs(self) -> List[OutputParam]:
        return [
            OutputParam(
                "f",
                type_hint=dict,
                description="Foundry feature dict consumed by the token initializer",
            ),
            OutputParam(
                "coord_atom_lvl_to_be_noised",
                type_hint=torch.Tensor,
                description="Reference atom-level coordinates [D, L_atom, 3]",
            ),
            OutputParam(
                "motif_mask",
                type_hint=torch.Tensor,
                description="Atom-level boolean mask for motif (fixed) positions [L_atom]",
            ),
            OutputParam(
                "motif_token_mask",
                type_hint=torch.Tensor,
                description="Residue-level boolean mask for motif (fixed) positions [L]",
            ),
            OutputParam(
                "L",
                type_hint=int,
                description="Total length of the protein being designed, in residues",
            ),
            OutputParam(
                "batch_size",
                type_hint=int,
                description="Batch size (typically 1 for RFDiffusion)",
            ),
        ]

    def check_inputs(self, components, block_state):
        if block_state.contigs is None:
            raise ValueError("`contigs` must be provided to specify protein design regions")

    @torch.no_grad()
    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        self.check_inputs(components, block_state)

        contigs = block_state.contigs
        input_xyz = block_state.input_xyz

        if isinstance(contigs, list):
            contig_str = "/".join(contigs)
        else:
            contig_str = contigs

        L, motif_ranges = parse_contig_string(contig_str)

        # Motif conditioning needs a reference structure so the transform pipeline can build
        # per-atom features for the fixed residues. A coordinate tensor alone cannot supply the
        # element, atom-name and occupancy annotations those features are derived from.
        if motif_ranges:
            raise ValueError(
                f"Motif-conditioned contigs are not supported yet, got `contigs={contig_str!r}`. "
                "Pass a plain design length such as `contigs='100'`."
            )
        if input_xyz is not None:
            raise ValueError(
                "`input_xyz` is not supported yet. Pass a plain design length such as "
                "`contigs='100'`."
            )

        batch_size = 1
        example = build_design_features(
            length=L,
            diffusion_batch_size=batch_size,
            sigma_data=components.scheduler.config.sigma_data,
        )

        block_state.f = example["feats"]
        block_state.coord_atom_lvl_to_be_noised = example["coord_atom_lvl_to_be_noised"]
        block_state.motif_mask = example["feats"]["is_motif_atom_with_fixed_coord"]
        block_state.motif_token_mask = example["feats"]["is_motif_token_with_fully_fixed_coord"]
        block_state.L = L
        block_state.batch_size = batch_size

        self.set_block_state(state, block_state)
        return components, state


class RFDiffusionSetTimestepsStep(ModularPipelineBlocks):
    """
    Set up the EDM noise schedule for RFDiffusion3.
    """

    model_name = "rfdiffusion"

    @property
    def description(self) -> str:
        return "Sets up the EDM noise schedule matching the original inference sampler."

    @property
    def expected_components(self) -> List[ComponentSpec]:
        return [
            ComponentSpec("scheduler", description="RFDiffusion3 EDM scheduler"),
        ]

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam(
                "num_inference_steps",
                default=None,
                type_hint=int,
                description="Number of denoising steps (default: use scheduler config)",
            ),
            InputParam("L", required=True, type_hint=int, description="Protein length"),
        ]

    @property
    def intermediate_outputs(self) -> List[OutputParam]:
        return [
            OutputParam(
                "noise_schedule",
                type_hint=torch.Tensor,
                description="EDM noise schedule [num_timesteps] from high to low noise",
            ),
            OutputParam(
                "num_inference_steps",
                type_hint=int,
                description="Number of inference steps",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        # A linear stand-in for the EDM schedule silently changes the sampler, so require
        # the real one rather than degrading.
        if components.scheduler is None:
            raise ValueError(
                "`scheduler` is not loaded. Call `load_components(trust_remote_code=True)` on the "
                "pipeline before calling it."
            )

        noise_schedule = components.scheduler.get_noise_schedule()

        block_state.noise_schedule = noise_schedule
        block_state.num_inference_steps = len(noise_schedule)

        self.set_block_state(state, block_state)
        return components, state


class RFDiffusionPrepareLatentsStep(ModularPipelineBlocks):
    """
    Prepare initial noised coordinates for RFDiffusion3.

    Matches the original _get_initial_structure:
        noise = c0 * randn(D, L, 3)
        noise[..., is_motif, :] = 0
        X_L = noise + coord_motif
    """

    model_name = "rfdiffusion"

    @property
    def description(self) -> str:
        return (
            "Prepares initial coordinates by sampling Gaussian noise scaled by "
            "the first noise schedule value, matching the original sampler."
        )

    @property
    def expected_components(self) -> List[ComponentSpec]:
        return [
            ComponentSpec("scheduler", description="RFDiffusion3 EDM scheduler"),
            ComponentSpec("transformer", description="RFDiffusion transformer model"),
        ]

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam("generator", type_hint=torch.Generator, description="Random generator for reproducibility"),
            InputParam("diffusion_batch_size", default=1, type_hint=int, description="Number of samples to generate in parallel"),
            InputParam("L", required=True, type_hint=int, description="Protein length"),
            InputParam("f", required=True, type_hint=dict),
            InputParam("coord_atom_lvl_to_be_noised", required=True, type_hint=torch.Tensor),
            InputParam("motif_mask", required=True, type_hint=torch.Tensor),
            InputParam("noise_schedule", required=True, type_hint=torch.Tensor),
        ]

    @property
    def intermediate_outputs(self) -> List[OutputParam]:
        return [
            OutputParam("xyz", type_hint=torch.Tensor, description="Initial noised coordinates [D, L_atom, 3]"),
            OutputParam("initializer_outputs", type_hint=dict, description="Embedded conditioning, reused every step"),
        ]

    @torch.no_grad()
    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        noise_schedule = block_state.noise_schedule
        generator = block_state.generator
        D = block_state.diffusion_batch_size or 1
        device = components.transformer.device

        # The feature dict is built on CPU and is read on every denoising step, so move it once.
        f = {k: v.to(device) if torch.is_tensor(v) else v for k, v in block_state.f.items()}
        coord = block_state.coord_atom_lvl_to_be_noised.to(device)
        motif_mask = f["is_motif_atom_with_fixed_coord"]

        # Matches rfd3.model.inference_sampler._get_initial_structure:
        #   noise = c0 * randn(D, L, 3); noise[..., is_motif, :] = 0; X_L = noise + coord
        c0 = noise_schedule[0].to(device)
        L_atom = coord.shape[-2]
        noise = c0 * torch.randn((D, L_atom, 3), device=device, generator=generator)
        noise[..., motif_mask, :] = 0.0
        xyz = noise + coord

        block_state.f = f
        block_state.xyz = xyz
        block_state.initializer_outputs = components.transformer.encode_conditioning(f)

        self.set_block_state(state, block_state)
        return components, state