"""Detect cancerous regions in a whole slide image."""

from datetime import datetime
import getpass
import json
import os
from pathlib import Path
import platform
import shutil
import subprocess
import sys
import typing

import click

from .._modellib.run_inference import run_inference
from .._modellib import models
from .._patchlib.create_dense_patch_grid import create_grid_and_save_multi_slides
from .._patchlib.create_patches_fp import create_patches

PathType = typing.Union[str, Path]


def _num_cpus() -> int:
    """Get number of CPUs on the system."""
    try:
        return len(os.sched_getaffinity(0))
    # os.sched_getaffinity seems to be linux only.
    except AttributeError:
        count = os.cpu_count()  # potentially None
        return count or 0


def _inside_container() -> str:
    if Path("/.dockerenv").exists():
        return "yes, docker"
    elif (
        Path("/singularity").exists()
        or Path("/singularity.d").exists()
        or Path("/.singularity.d").exists()
    ):
        # TODO: apptainer might change the name of this directory.
        return "yes, apptainer/singularity"
    return "no"


def _get_timestamp() -> str:
    dt = datetime.now().astimezone()
    # Thu Aug 25 23:32:17 2022 EDT
    return dt.strftime("%c %Z")


def _print_system_info() -> None:
    """Print information about the system."""
    import torch
    import torchvision
    from .. import __version__

    click.secho(f"\nRunning wsinfer version {__version__}", fg="green")
    print("\nIf you run into issues, please submit a new issue at")
    print("https://github.com/kaczmarj/patch-classification-pipeline/issues/new")
    print("\nSystem information")
    print("------------------")
    print(f"Timestamp: {_get_timestamp()}")
    print(f"{platform.platform()}")
    try:
        print(f"User: {getpass.getuser()}")
    except KeyError:
        # If /etc/passwd does not include the username of the current user ID, a
        # KeyError is thrown. This could happen in a Docker image running as a different
        # user with `--user $(id -u):$(id -g)` but that does not bind mount the
        # /etc/passwd file.
        print("User: UNKNOWN")
    print(f"Hostname: {platform.node()}")
    print(f"Working directory: {os.getcwd()}")
    print(f"In container: {_inside_container()}")
    print(f"Python executable: {sys.executable}")
    print(f"Python version: {platform.python_version()}")
    print(f"  Torch version: {torch.__version__}")
    print(f"  Torchvision version: {torchvision.__version__}")
    cuda_is_available = torch.cuda.is_available()
    if cuda_is_available:
        click.secho("GPU available", fg="green")
        click.secho(f"  Using {torch.cuda.device_count()} GPU(s)", fg="green")
        cuda_ver = torch.version.cuda or "NOT FOUND"
        print(f"  CUDA version: {cuda_ver}")
    else:
        click.secho("GPU not available", bg="red", fg="black")
    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "NOT SET")
    print(f"CUDA_VISIBLE_DEVICES: {cuda_visible_devices}")
    if torch.version.cuda is None:
        click.secho("\n*******************************************", fg="yellow")
        click.secho("GPU WILL NOT BE USED", fg="yellow")
        if torch.version.cuda is None:
            click.secho("  CPU-only version of PyTorch is installed", fg="yellow")
        click.secho("*******************************************", fg="yellow")
    elif not cuda_is_available:
        click.secho("\n*******************************************", fg="yellow")
        click.secho("GPU WILL NOT BE USED", fg="yellow")
        if torch.version.cuda is None:
            click.secho("  CUDA DEVICES NOT AVAILABLE", fg="yellow")
        click.secho("*******************************************", fg="yellow")


def _get_info_for_save(weights: models.Weights):
    """Get dictionary with information about the run. To save as JSON in output dir."""

    import torch
    from .. import __version__

    here = Path(__file__).parent.resolve()

    def get_git_info():
        here = Path(__file__).parent.resolve()

        def get_stdout(args) -> str:
            proc = subprocess.run(args, capture_output=True, cwd=here)
            return "" if proc.returncode != 0 else proc.stdout.decode().strip()

        git_remote = get_stdout("git config --get remote.origin.url".split())
        git_branch = get_stdout("git rev-parse --abbrev-ref HEAD".split())
        git_commit = get_stdout("git rev-parse HEAD".split())

        # https://stackoverflow.com/a/3879077/5666087
        cmd = subprocess.run("git diff-index --quiet HEAD --".split(), cwd=here)
        uncommitted_changes = cmd.returncode != 0
        return {
            "git_remote": git_remote,
            "git_branch": git_branch,
            "git_commit": git_commit,
            "uncommitted_changes": uncommitted_changes,
        }

    # Test if we are in a git repo. If we are, then get git info.
    git_program = shutil.which("git")
    git_installed = git_program is not None
    is_git_repo = False
    if git_installed:
        cmd = subprocess.run(
            [str(git_program), "branch"],
            cwd=here,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        is_git_repo = cmd.returncode == 0
    git_info = None
    if git_installed and is_git_repo:
        git_info = get_git_info()

    weights_file = weights.file
    if weights_file is None:
        if weights.url_file_name is None:
            raise TypeError("url_file_name must not be None if file is None.")
        weights_file = str(
            Path(torch.hub.get_dir()) / "checkpoints" / weights.url_file_name
        )

    return {
        "model_weights": {
            "name": weights.name,
            "architecture": weights.architecture,
            "weights_url": weights.url,
            "weights_url_file_name": weights.url_file_name,
            "weights_file": weights_file,
            "weights_sha256": weights.get_sha256_of_weights(),
            "class_names": weights.class_names,
            "num_classes": weights.num_classes,
            "patch_size_pixels": weights.patch_size_pixels,
            "spacing_um_px": weights.spacing_um_px,
            "transform": {
                "resize_size": weights.transform.resize_size,
                "mean": weights.transform.mean,
                "std": weights.transform.std,
            },
            "metadata": weights.metadata or None,
        },
        "runtime": {
            "version": __version__,
            "working_dir": os.getcwd(),
            "args": " ".join(sys.argv),
            "python_executable": sys.executable,
            "python_version": platform.python_version(),
            "in_container": _inside_container(),
            "pytorch_version": torch.__version__,
            "cuda_version": torch.version.cuda,
            "git": git_info,
        },
        "timestamp": _get_timestamp(),
    }


@click.command(context_settings=dict(auto_envvar_prefix="WSINFER"))
@click.pass_context
@click.option(
    "--wsi-dir",
    type=click.Path(exists=True, file_okay=False, path_type=Path, resolve_path=True),
    required=True,
    help="Directory containing whole slide images. This directory can *only* contain"
    " whole slide images.",
)
@click.option(
    "--results-dir",
    type=click.Path(file_okay=False, path_type=Path, resolve_path=True),
    required=True,
    help="Directory to store results. If directory exists, will skip"
    " whole slides for which outputs exist.",
)
@click.option(
    "--model",
    type=click.Choice(sorted({a for a, _ in models.list_all_models_and_weights()})),
    help="Model architecture to use. Not required if 'config' is used.",
)
@click.option(
    "--weights",
    type=click.Choice(sorted({w for _, w in models.list_all_models_and_weights()})),
    help="Name of weights to use for the model. Not required if 'config' is used.",
)
@click.option(
    "--config",
    type=click.Path(exists=True, dir_okay=False, path_type=Path, resolve_path=True),
    help=(
        "Path to configuration for architecture and weights. Use this option if the"
        " model weights are not registered in wsinfer. Mutually exclusive with"
        " 'model' and 'weights'."
    ),
)
@click.option(
    "--batch-size",
    type=click.IntRange(min=1),
    default=32,
    show_default=True,
    help="Batch size during model inference. If using multiple GPUs, increase the"
    " batch size.",
)
@click.option(
    "--num-workers",
    default=min(_num_cpus(), 8),  # Use at most 8 workers by default.
    show_default=True,
    type=click.IntRange(min=0),
    help="Number of workers to use for data loading during model inference (n=0 for"
    " single thread). Set this to the number of cores on your machine or lower.",
)
@click.option(
    "--speedup/--no-speedup",
    default=False,
    show_default=True,
    help="JIT-compile the model for potential speedups.",
)
@click.option(
    "--dense-grid/--no-dense-grid",
    default=False,
    show_default=True,
    help="Use a dense grid of patch coordinates. Patches will be present even if no"
    " tissue is present",
)
def run(
    ctx: click.Context,
    *,
    wsi_dir: Path,
    results_dir: Path,
    model: typing.Optional[str],
    weights: typing.Optional[str],
    config: typing.Optional[Path],
    batch_size: int,
    num_workers: int = 0,
    speedup: bool = False,
    dense_grid: bool = False,
):
    """Run model inference on a directory of whole slide images.

    This command will create a tissue mask of each WSI. Then patch coordinates will be
    computed. The chosen model will be applied to each patch, and the results will be
    saved to a CSV in `RESULTS_DIR/model-output`.

    Example:

    CUDA_VISIBLE_DEVICES=0 wsinfer run --wsi_dir slides/ --results_dir results
    --model resnet34 --weights TCGA-BRCA-v1 --batch_size 32 --num_workers 4

    To list all available models and weights, use `wsinfer list`.
    """
    if model is None and weights is None and config is None:
        raise click.UsageError("one of (model and weights) or config is required.")
    elif (model is not None or weights is not None) and config is not None:
        raise click.UsageError("model and weights are mutually exclusive with config.")
    elif (model is not None) ^ (weights is not None):  # XOR
        raise click.UsageError("model and weights must both be set if one is set.")

    wsi_dir = wsi_dir.resolve()
    results_dir = results_dir.resolve()

    if not wsi_dir.exists():
        raise FileNotFoundError(f"Whole slide image directory not found: {wsi_dir}")

    # Test that wsi dir actually includes files. This is here for an interesting edge
    # case. When using a Linux container and if the data directory is symlinked from a
    # different directory, both directories need to be bind mounted onto the container.
    # If only the symlinked directory is included, then the patching script will fail,
    # even though it looks like there are files in the wsi_dir directory.
    files_in_wsi_dir = [p for p in wsi_dir.glob("*") if p.exists()]
    if not files_in_wsi_dir:
        raise FileNotFoundError(f"no files exist in the slide directory: {wsi_dir}")

    _print_system_info()

    print("\nCommand line arguments")
    print("----------------------")
    for key, value in ctx.params.items():
        print(f"{key} = {value}")
    print("----------------------\n")

    # Get weights object before running the patching script because we need to get the
    # necessary spacing and patch size.
    if model is not None and weights is not None:
        weights_obj = models.get_model_weights(model, name=weights)
    elif config is not None:
        weights_obj = models.Weights.from_yaml(config)

    click.secho("\nFinding patch coordinates...\n", fg="green")
    if dense_grid:
        click.echo("Not using a tissue mask.")
        create_grid_and_save_multi_slides(
            wsi_dir=wsi_dir,
            results_dir=results_dir,
            orig_patch_size=weights_obj.patch_size_pixels,
            patch_spacing_um_px=weights_obj.spacing_um_px,
        )
    else:
        create_patches(
            source=str(wsi_dir),
            save_dir=str(results_dir),
            patch_size=weights_obj.patch_size_pixels,
            patch_spacing=weights_obj.spacing_um_px,
            seg=True,
            patch=True,
            preset="tcga.csv",
        )

    click.secho("\nRunning model inference.\n", fg="green")
    failed_patching, failed_inference = run_inference(
        wsi_dir=wsi_dir,
        results_dir=results_dir,
        weights=weights_obj,
        batch_size=batch_size,
        num_workers=num_workers,
        speedup=speedup,
    )

    if failed_patching:
        click.secho(f"\nPatching failed for {len(failed_patching)} slides", fg="yellow")
        click.secho("\n".join(failed_patching), fg="yellow")
    if failed_inference:
        click.secho(
            f"\nInference failed for {len(failed_inference)} slides", fg="yellow"
        )
        click.secho("\n".join(failed_inference), fg="yellow")

    timestamp = datetime.now().astimezone().strftime("%Y%m%dT%H%M%S")
    run_metadata_outpath = results_dir / f"run_metadata_{timestamp}.json"
    click.echo(f"Saving metadata about run to {run_metadata_outpath}")
    run_metadata = _get_info_for_save(weights_obj)
    with open(run_metadata_outpath, "w") as f:
        json.dump(run_metadata, f, indent=2)

    click.secho("Finished.", fg="green")
