import multiprocessing as mp
from pathlib import Path

import click
import numpy as np
import pandas as pd
from tqdm import tqdm

from openfold3.core.data.io.structure.atom_array import read_atomarray_from_npz
from openfold3.core.data.resources.residues import MoleculeType


@click.command()
@click.option(
    "--structure_directory",
    help=(
        "Directory containing cif files to count unresolved residues from. "
        "Mutually exclusive with --structure_array_directory."
    ),
    type=click.Path(
        file_okay=False,
        dir_okay=True,
        path_type=Path,
    ),
    default=None,
)
@click.option(
    "--structure_array_directory",
    help=(
        "Directory containing subdirectories of preprocessed cif files to count "
        "unresolved residues from generated by "
        "`scripts/preprocess_template_structures_of3.py`. Mutually exclusive with "
        "--structure_directory."
    ),
    type=click.Path(
        file_okay=False,
        dir_okay=True,
        path_type=Path,
    ),
    default=None,
)
@click.option(
    "--output_directory",
    help="Directory to write output files to.",
    type=click.Path(
        file_okay=False,
        dir_okay=True,
        path_type=Path,
    ),
    required=True,
)
@click.option(
    "--moltypes",
    help=(
        "Comma-separated list of molecule types to consider when counting unresolved "
        "residues. Options are PROTEIN, DNA, RNA, LIGAND. Default: PROTEIN"
    ),
    type=str,
    default="PROTEIN",
)
@click.option(
    "--min_n_res_unresolved",
    help=(
        "Minimum number of unresolved residues required to include entry in output "
        "list. Default: 40"
    ),
    type=int,
    default=40,
)
@click.option(
    "--num_workers",
    required=True,
    type=int,
    help=(
        "Number of workers to parallelize the template cache computation and filtering"
        " over."
    ),
)
@click.option(
    "--chunksize",
    type=int,
    help=("Chunk size to use for distributing work among workers."),
    default=1,
)
def main(
    structure_directory: Path | None,
    structure_array_directory: Path | None,
    output_directory: Path,
    moltypes: str,
    min_n_res_unresolved: int,
    num_workers: int,
    chunksize: int,
):
    moltypes = [
        MoleculeType[moltype_str.strip().upper()] for moltype_str in moltypes.split(",")
    ]
    output_directory.mkdir(parents=True, exist_ok=True)
    unresolved_counter = UnresolvedResidueCounter(
        structure_directory=structure_directory,
        structure_array_directory=structure_array_directory,
        output_directory=output_directory,
        moltypes=moltypes,
        min_n_res_unresolved=min_n_res_unresolved,
        num_workers=num_workers,
        chunksize=chunksize,
    )
    unresolved_counter()


class UnresolvedResidueCounter:
    """Unresolved residue counter pipeline."""

    def __init__(
        self,
        structure_directory: Path | None,
        structure_array_directory: Path | None,
        output_directory: Path,
        moltypes: list[MoleculeType],
        min_n_res_unresolved: int,
        num_workers: int,
        chunksize: int,
    ) -> None:
        if (
            (structure_directory is not None)
            and (structure_array_directory is not None)
        ) or ((structure_directory is None) and (structure_array_directory is None)):
            raise ValueError(
                "Exactly one of --structure_directory or --structure_array_directory "
                "must be provided."
            )

        self.structure_directory = structure_directory
        self.structure_array_directory = structure_array_directory
        self.output_directory = output_directory
        self.moltypes = moltypes
        self.min_n_res_unresolved = min_n_res_unresolved
        self.num_workers = num_workers
        self.chunksize = chunksize

        self.entry_ids = []
        self.n_res_unresolved_counts = []

        if self.structure_array_directory is not None:
            self.structure_files = None
            self.structure_array_subdirs = list(
                [f for f in structure_array_directory.iterdir()]
            )
            print(
                f"Found {len(self.structure_array_subdirs)} structure array "
                "subdirectories."
            )
        else:
            self.structure_files = list([f for f in structure_directory.iterdir()])
            self.structure_array_subdirs = None
            print(f"Found {len(self.structure_files)} structure files.")

    def count_unresolved_in_cif(self) -> tuple[str, int]:
        "Counts unresolved residues in a single cif file."
        raise NotImplementedError(
            "Counting unresolved residues in cif file not yet implemented."
        )

    def process_structure_directory(self) -> None:
        raise NotImplementedError("Processing structure directory not yet implemented.")

    def count_unresolved_in_atom_arrays(self, subdir: Path) -> tuple[str, int]:
        "Counts unresolved residues in a set of per-chain Atom Array npz files."
        n_res_unresolved_total = 0
        for chain_array_file in subdir.iterdir():
            if chain_array_file.stem != "chain_id_to_moltype":
                atom_array = read_atomarray_from_npz(chain_array_file)
                moltype_id = np.unique(atom_array.molecule_type_id)
                if len(moltype_id) != 1:
                    print(
                        f"Found multiple molecule types in {chain_array_file.stem}, "
                        "skipping..."
                    )
                else:
                    if MoleculeType(moltype_id.item()) in self.moltypes:
                        n_res_unresolved_total += len(
                            np.unique(atom_array[atom_array.occupancy == 0].res_id)
                        )

        entry_id = subdir.stem

        return entry_id, n_res_unresolved_total

    def process_structure_array_directory(self) -> None:
        with mp.Pool(self.num_workers) as pool:
            for entry_id, n_res_unresolved_total in tqdm(
                pool.imap_unordered(
                    self.count_unresolved_in_atom_arrays,
                    self.structure_array_subdirs,
                    chunksize=self.chunksize,
                ),
                total=len(self.structure_array_subdirs),
                desc=(
                    "Counting unresolved residues in chains with molecule "
                    f"type(s) {self.moltypes}"
                ),
            ):
                self.entry_ids.append(entry_id)
                self.n_res_unresolved_counts.append(n_res_unresolved_total)

    def save_results(self) -> None:
        df = pd.DataFrame(
            {
                "entry_id": self.entry_ids,
                "n_res_unresolved": self.n_res_unresolved_counts,
            }
        )
        df.to_csv(
            self.output_directory / "n_res_unresolved_counts.tsv",
            index=False,
            header=True,
            sep="\t",
        )
        df_unresolved = df[df["n_res_unresolved"] >= self.min_n_res_unresolved]
        output_file_unresolved = self.output_directory / "entries_unresolved.tsv"
        df_unresolved.to_csv(output_file_unresolved, index=False, header=True, sep="\t")
        print(f"Saved unresolved residue counts to {self.output_directory}")

    def __call__(self) -> None:
        if self.structure_directory is not None:
            self.process_structure_directory()
        else:
            self.process_structure_array_directory()

        self.save_results()

        print("Done.")


if __name__ == "__main__":
    main()
