from dataclasses import dataclass
from pathlib import Path
import re
from typing import Dict, List, Optional, Union

import numpy as np
import pytorch_lightning as pl
import torch
from torch import Tensor
from torch.utils.data import DataLoader
from rdkit.Chem import Mol
from boltzgen.data import const
from boltzgen.data.data import Input
from boltzgen.data.feature.featurizer import Featurizer
from boltzgen.data.pad import pad_to_max
from boltzgen.data.mol import load_canonicals, load_molecules
from boltzgen.data.parse.schema import YamlDesignParser
from boltzgen.data.template.features import load_dummy_templates
from boltzgen.data.tokenize.tokenizer import Tokenizer
from boltzgen.data.data import Input
from boltzgen.data.select.protein import ProteinSelector


@dataclass
class DataConfig:
    """Data configuration."""

    moldir: str
    multiplicity: int
    yaml_path: Union[List[str], str]
    tokenizer: Tokenizer
    featurizer: Featurizer
    backbone_only: bool = False
    atom14: bool = False
    atom37: bool = False
    design: bool = True
    compute_affinity: bool = False
    disulfide_prob: float = 1.0
    disulfide_on: bool = False
    skip_existing: bool = False
    skip_offset: int = 0
    diffusion_samples: int = 1
    output_dir: Optional[str] = None
  
   


@dataclass
class Dataset:
    yaml_path: Union[List[str], str]
    tokenizer: Tokenizer
    featurizer: Featurizer
    multiplicity: int = 1


def collate(data: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
    """Collate the data.

    Parameters
    ----------
    data : List[Dict[str, Tensor]]
        The data to collate.

    Returns
    -------
    Dict[str, Tensor]
        The collated data.

    """
    # Get the keys
    keys = data[0].keys()

    # Collate the data
    collated = {}
    for key in keys:
        values = [d[key] for d in data]

        if key not in [
            "all_coords",
            "all_resolved_mask",
            "crop_to_all_atom_map",
            "chain_symmetries",
            "amino_acids_symmetries",
            "ligand_symmetries",
            "activity_name",
            "activity_qualifier",
            "sid",
            "cid",
            "aid",
            "normalized_protein_accession",
            "pair_id",
            "record",
            "id",
            "structure",
            "tokenized",
            "structure_bonds",
            "extra_mols",
            "data_sample_idx",
        ]:
            # Check if all have the same shape
            shape = values[0].shape
            if not all(v.shape == shape for v in values):
                values = pad_to_max(values, 0)
            else:
                values = torch.stack(values, dim=0)

        # Stack the values
        collated[key] = values

    return collated


class PredictionDataset(torch.utils.data.Dataset):
    """Base iterable dataset."""

    def __init__(
        self,
        dataset: Dataset,
        canonicals: dict[str, Mol],
        moldir: str,
        backbone_only: bool = False,
        atom14: bool = False,
        atom37: bool = False,
        extra_features: Optional[List[str]] = None,
        design: bool = True,
        compute_affinity: bool = False,
        disulfide_prob: float = 1.0,
        disulfide_on: bool = False,
        skip_offset: int = 0,
    ) -> None:
        """Initialize the training dataset.

        Parameters
        ----------
        datasets : List[Dataset]
            The datasets to sample from.

        """
        super().__init__()
        self.dataset = dataset
        self.moldir = moldir
        self.canonicals = canonicals
        self.backbone_only = backbone_only
        self.atom14 = atom14
        self.atom37 = atom37
        self.skip_offset = skip_offset
        path = dataset.yaml_path
        self.yaml_paths = [path] if isinstance(path, str) else path

        for path in self.yaml_paths:
            filename = Path(path).name
            if re.search(r"_\d+\.yaml$", filename):
                raise ValueError(
                    f"Illegal YAML filename for '{str(path)}': names must not end with the pattern _\\d+\\.yaml so, e.g., the ends '_001.yaml' or '_4.yaml' are not allowed."
                    "This pattern is reserved for internal file indexing. Sorry :)"
                )
            if "_native" in filename:
                raise ValueError(
                    f"Illegal YAML filename for '{str(path)}': names must not contain '_native' because this substring is reserved for native structure companions. Sorry :)"
                )
        self.extra_features = (
            set(extra_features) if extra_features is not None else set()
        )
        self.selector = (
            ProteinSelector(  
                design_neighborhood_sizes=[2, 4, 6, 8, 10, 12, 14, 16, 18],
                substructure_neighborhood_sizes=[2, 4, 6, 8, 10, 12, 24],
                structure_condition_prob=1.0,
                distance_noise_std=1,
                run_selection=True,
                specify_binding_sites=True,
                ss_condition_prob=0.1,
                select_all=False,
                chain_reindexing=False,
            )
        )
        self.design = design
        self.compute_affinity = compute_affinity
        self.disulfide_prob = disulfide_prob
        self.disulfide_on = disulfide_on

        self.mols = {}
        self.parser = YamlDesignParser(mol_dir=self.moldir)

    def __getitem__(self, idx: int) -> Dict:
        """Get an item from the dataset.

        Returns
        -------
        Dict[str, Tensor]
            The sampled data features.

        """
        path = Path(self.yaml_paths[idx % len(self.yaml_paths)])
        feat = self.get_sample(path)
        data_sample_idx = idx // len(self.yaml_paths) + self.skip_offset
        if self.dataset.multiplicity > 1:
            feat["data_sample_idx"] = data_sample_idx
        return feat

    def get_sample(self, path: Path, sample_id: Optional[str] = None) -> Dict:
        # Get itemn also needs to take a smaple id as input
        parsed = self.parser.parse_yaml(
            path, mol_dir=self.moldir, mols=self.mols
        )
        structure = parsed.structure
        design_info = parsed.design_info

        # Tokenize structure
        tokenized = self.dataset.tokenizer.tokenize(structure)

        # Transfer conditioning information that is stored in tokens
        token_to_res = tokenized.token_to_res
        tokenized.tokens["design_mask"] = design_info.res_design_mask[token_to_res]
        tokenized.tokens["binding_type"] = design_info.res_binding_type[token_to_res]
        tokenized.tokens["structure_group"] = design_info.res_structure_groups[
            token_to_res
        ]

        # Propagate design mask to obtain chain_design_mask (True whenever something is covalently bound to any residue that is in a chain that contains a design residue).
        chain_design_mask = tokenized.tokens["design_mask"].astype(bool)
        asym_id = tokenized.tokens["asym_id"]
        while True:
            design_chains = np.unique(asym_id[chain_design_mask])
            chain_propagated = np.isin(asym_id, design_chains)
            for i, j, _ in tokenized.bonds:
                if any([chain_propagated[i], chain_propagated[j]]):
                    chain_propagated[i] = True
                    chain_propagated[j] = True
            if np.equal(chain_propagated, chain_design_mask).all():
                break
            chain_design_mask = chain_propagated.astype(bool)

        # Try to find molecules in the dataset moldir if provided
        # Find missing ones in global moldir and check if all found
        molecules = {}
        molecules.update(self.canonicals)
        mol_names = set(tokenized.tokens["res_name"].tolist())
        mol_names = mol_names - set(self.canonicals.keys())
        mol_names = mol_names - set(parsed.extra_mols.keys())
        if self.moldir is not None:
            molecules.update(load_molecules(self.moldir, mol_names))

        mol_names = mol_names - set(molecules.keys())
        molecules.update(load_molecules(self.moldir, mol_names))
        molecules.update(parsed.extra_mols)

        # Finalize input data
        input_data = Input(
            tokens=tokenized.tokens,
            bonds=tokenized.bonds,
            token_to_res=token_to_res,
            structure=structure,
            msa={},
            templates=None,
        )

        # Compute features
        features = self.dataset.featurizer.process(
            input_data,
            molecules=molecules,
            random=np.random.default_rng(None),
            training=False,
            max_seqs=1,
            backbone_only=self.backbone_only,
            atom14=self.atom14,
            atom37=self.atom37,
            design=self.design,
            override_method="X-RAY DIFFRACTION",
            compute_affinity=self.compute_affinity,
            disulfide_prob=self.disulfide_prob,
            disulfide_on=self.disulfide_on,
        )

        # transfer secondary structure conditioning
        ss_type = design_info.res_ss_types[token_to_res]
        features["ss_type"] = torch.from_numpy(ss_type).to(features["ss_type"])
        features["design_ss_mask"][ss_type != const.ss_type_ids["UNSPECIFIED"]] = 1

        # set chain_design_mask
        features["chain_design_mask"] = torch.from_numpy(chain_design_mask)

        # Compute template features
        templates_features = load_dummy_templates(
            tdim=1, num_tokens=len(features["res_type"])
        )
        features.update(templates_features)

        # set last necessary features
        features["idx_dataset"] = torch.tensor(1)

        # If a smaple id is provided then this should be the sample id instead of path.stem
        if sample_id is not None:
            features["id"] = sample_id
        else:
            features["id"] = path.stem
        if "structure" in self.extra_features:
            features["structure"] = structure
        if "tokenized" in self.extra_features:
            features["tokenized"] = tokenized

        return features

    def __len__(self) -> int:
        """Get the length of the dataset.

        Returns
        -------
        int
            The length of the dataset.

        """
        total = len(self.yaml_paths) * (self.dataset.multiplicity - self.skip_offset)
        return max(total, 0)


class FromYamlDataModule(pl.LightningDataModule):
    """DataModule for BoltzGen."""

    def __init__(
        self, cfg: DataConfig, batch_size, num_workers, pin_memory, extra_features=None
    ) -> None:
        """Initialize the DataModule.

        Parameters
        ----------
        config : DataConfig
            The data configuration.

        """
        super().__init__()

        if cfg.skip_existing and cfg.output_dir is not None:
            design_dir = Path(cfg.output_dir)
            max_idx: int = -1
            if design_dir.exists():
                pattern = re.compile(r"_(\d+)(?:\.[^.]+)$")
                max_idx = max(
                    (
                        int(m.group(1))
                        for fp in design_dir.iterdir()
                        if fp.suffix in {".cif", ".pdb"}
                        and not any(s in fp.name for s in ("_native.cif", "_metadata.npz"))
                        for m in [pattern.search(fp.name)]
                        if m
                    ),
                    default=-1,
                )
            n_samples = getattr(cfg, "diffusion_samples", 1)
            cfg.skip_offset = (max_idx // max(n_samples, 1)) + 1 if max_idx >= 0 else 0
            
        self.cfg = cfg
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.pin_memory = pin_memory
        self.collate = collate

        dataset = Dataset(
            yaml_path=cfg.yaml_path,
            multiplicity=cfg.multiplicity,
            tokenizer=cfg.tokenizer,
            featurizer=cfg.featurizer,
        )

        # Load canonical molecules
        canonicals = load_canonicals(cfg.moldir)

        self.predict_set = PredictionDataset(
            dataset=dataset,
            canonicals=canonicals,
            moldir=Path(cfg.moldir),
            backbone_only=cfg.backbone_only,
            atom14=cfg.atom14,
            atom37=cfg.atom37,
            extra_features=extra_features,
            design=cfg.design,
            compute_affinity=cfg.compute_affinity,
            disulfide_prob=cfg.disulfide_prob,
            disulfide_on=cfg.disulfide_on,
            skip_offset=cfg.skip_offset,
        )

    def predict_dataloader(self) -> DataLoader:
        """Get the training dataloader.

        Returns
        -------
        DataLoader
            The training dataloader.

        """
        return DataLoader(
            self.predict_set,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=self.pin_memory,
            shuffle=False,
            collate_fn=collate,
        )

    def transfer_batch_to_device(
        self,
        batch: Dict,
        device: torch.device,
        dataloader_idx: int = 0,  # noqa: ARG002
    ) -> Dict:
        """Transfer a batch to the given device.

        Parameters
        ----------
        batch : Dict
            The batch to transfer.
        device : torch.device
            The device to transfer to.

        Returns
        -------
        np.Any
            The transferred batch.

        """
        for key in batch:
            if key not in [
                "all_coords",
                "all_resolved_mask",
                "crop_to_all_atom_map",
                "chain_symmetries",
                "amino_acids_symmetries",
                "ligand_symmetries",
                "activity_name",
                "activity_qualifier",
                "sid",
                "cid",
                "normalized_protein_accession",
                "pair_id",
                "record",
                "id",
                "structure",
                "tokenized",
                "structure_bonds",
                "extra_mols",
                "data_sample_idx",
            ]:
                batch[key] = batch[key].to(device)
        return batch
