import argparse
import os
from typing import Literal

import numpy as np
import torch
from monai.data import PersistentDataset, load_decathlon_datalist
from monai.transforms import (
    Compose,
    ConcatItemsd,
    DeleteItemsd,
    EnsureTyped,
    LoadImaged,
    RandCropByPosNegLabeld,
    RandRotate90d,
    ToTensord,
    Transform,
    Transposed,
)
from torch.utils.data.dataloader import default_collate

from .custom_transforms import (
    ClipMaskIntensityPercentilesd,
    ElementwiseProductd,
    NormalizeIntensity_customd,
    NormalizePSAd,
)


def list_data_collate(batch: list):
    """
    Combine instances from a list of dicts into a single dict, by stacking them along first dim
    [{'image' : 3xHxW}, {'image' : 3xHxW}, {'image' : 3xHxW}...] - > {'image' : Nx3xHxW}
    followed by the default collate which will form a batch BxNx3xHxW
    """

    for i, item in enumerate(batch):
        data = item[0]
        data["image"] = torch.stack([ix["image"] for ix in item], dim=0)

        if all("final_heatmap" in ix for ix in item):
            data["final_heatmap"] = torch.stack([ix["final_heatmap"] for ix in item], dim=0)
            data["smooth_mask"] = torch.stack([ix["smooth_mask"] for ix in item], dim=0)

        batch[i] = data
    return default_collate(batch)


def data_transform(args: argparse.Namespace, split) -> Transform:
    if split == "train":
        transform = Compose(
            [
                LoadImaged(
                    keys=["image", "mask", "dwi", "adc", "heatmap", "smooth_mask"],
                    reader="ITKReader",
                    ensure_channel_first=True,
                    dtype=np.float32,
                ),
                # LabelEncodeIntegerGraded(keys=["label"], num_classes=args.num_classes),
                ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
                ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"),
                NormalizeIntensity_customd(keys=["image"], mask_key="mask"),
                NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"),
                ConcatItemsd(
                    keys=["image", "dwi", "adc"], name="image", dim=0
                ),  # stacks to (3, H, W)
                ElementwiseProductd(keys=["heatmap", "smooth_mask"], output_key="final_heatmap"),
                RandCropByPosNegLabeld(
                    keys=["image", "final_heatmap", "smooth_mask"],
                    label_key="smooth_mask",
                    spatial_size=(args.tile_size, args.tile_size, args.depth),
                    pos=1,
                    neg=0,
                    num_samples=args.tile_count,
                ),
                RandRotate90d(
                    keys=["image", "final_heatmap", "smooth_mask"],
                    prob=0.6,
                    spatial_axes=(0, 1),
                    max_k=3,
                ),
                NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std),
                EnsureTyped(keys=["label", "psa"], dtype=torch.float32),
                Transposed(keys=["image"], indices=(0, 3, 1, 2)),
                DeleteItemsd(keys=["dwi", "adc", "heatmap", "mask"]),
                ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask", "psa"]),
            ]
        )
    else:
        transform = Compose(
            [
                LoadImaged(
                    keys=["image", "mask", "dwi", "adc", "heatmap", "smooth_mask"],
                    reader="ITKReader",
                    ensure_channel_first=True,
                    dtype=np.float32,
                ),
                # LabelEncodeIntegerGraded(keys=["label"], num_classes=args.num_classes),
                ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
                ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"),
                NormalizeIntensity_customd(keys=["image"], mask_key="mask"),
                NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"),
                ConcatItemsd(
                    keys=["image", "dwi", "adc"], name="image", dim=0
                ),  # stacks to (3, H, W)
                ElementwiseProductd(keys=["heatmap", "smooth_mask"], output_key="final_heatmap"),
                RandCropByPosNegLabeld(
                    keys=["image", "final_heatmap", "smooth_mask"],
                    label_key="smooth_mask",
                    spatial_size=(args.tile_size, args.tile_size, args.depth),
                    pos=1,
                    neg=0,
                    num_samples=args.tile_count,
                ),
                NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std),
                EnsureTyped(keys=["label", "psa"], dtype=torch.float32),
                Transposed(keys=["image"], indices=(0, 3, 1, 2)),
                DeleteItemsd(keys=["dwi", "adc", "heatmap", "mask"]),
                ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask", "psa"]),
            ]
        )
    return transform


"""
def data_transform(args: argparse.Namespace) -> Transform:
    if args.use_heatmap:
        if args.use_psa:
            transform = Compose(
                [
                    LoadImaged(
                        keys=["image", "mask", "dwi", "adc", "heatmap","smooth_mask"],
                        reader="ITKReader",
                        ensure_channel_first=True,
                        dtype=np.float32,
                    ),
                    ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
                    ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"),
                    NormalizeIntensity_customd(keys=["image"], mask_key="mask"),
                    NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"),
                    ConcatItemsd(
                        keys=["image", "dwi", "adc"], name="image", dim=0
                    ),  # stacks to (3, H, W)
                    ElementwiseProductd(keys=["heatmap", "smooth_mask"], output_key="final_heatmap"),
                    RandCropByPosNegLabeld(
                        keys=["image", "final_heatmap", "smooth_mask"],
                        label_key="smooth_mask",
                        spatial_size=(args.tile_size, args.tile_size, args.depth),
                        pos=1,
                        neg=0,
                        num_samples=args.tile_count,
                    ),
                    NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std),
                    EnsureTyped(keys=["label", "psa"], dtype=torch.float32),
                    Transposed(keys=["image"], indices=(0, 3, 1, 2)),
                    DeleteItemsd(keys=[ "dwi", "adc", "heatmap", "mask"]),
                    ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask", "psa"]),
                ]
            )
        else:
            transform = Compose(
                [
                    LoadImaged(
                        keys=["image", "mask", "dwi", "adc", "heatmap","smooth_mask"],
                        reader="ITKReader",
                        ensure_channel_first=True,
                        dtype=np.float32,
                    ),
                    ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
                    ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"),
                    NormalizeIntensity_customd(keys=["image"], mask_key="mask"),
                    NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"),
                    ConcatItemsd(
                        keys=["image", "dwi", "adc"], name="image", dim=0
                    ),  # stacks to (3, H, W)
                    ElementwiseProductd(keys=["heatmap", "smooth_mask"], output_key="final_heatmap"),
                    #RandRotate90d(keys=["image", "final_heatmap", "smooth_mask"], prob=0.5, spatial_axes=(0, 1)),
                    RandCropByPosNegLabeld(
                        keys=["image", "final_heatmap", "smooth_mask"],
                        label_key="smooth_mask",
                        spatial_size=(args.tile_size, args.tile_size, args.depth),
                        pos=1,
                        neg=0,
                        num_samples=args.tile_count,
                    ),
                    EnsureTyped(keys=["label"], dtype=torch.float32),
                    Transposed(keys=["image"], indices=(0, 3, 1, 2)),
                    DeleteItemsd(keys=[ "dwi", "adc", "heatmap", "mask"]),
                    ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask"]),
                ]
            )
    else:
        if args.use_psa:
            transform = Compose(
                [
                    LoadImaged(
                        keys=["image", "mask", "dwi", "adc","smooth_mask"],
                        reader="ITKReader",
                        ensure_channel_first=True,
                        dtype=np.float32,
                    ),
                    ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
                    ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"),
                    NormalizeIntensity_customd(keys=["image"], mask_key="mask"),
                    NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"),
                    ConcatItemsd(
                        keys=["image", "dwi", "adc"], name="image", dim=0
                    ),  # stacks to (3, H, W)
                    RandCropByPosNegLabeld(
                        keys=["image", "smooth_mask"],
                        label_key="smooth_mask",
                        spatial_size=(args.tile_size, args.tile_size, args.depth),
                        pos=1,
                        neg=0,
                        num_samples=args.tile_count,
                    ),
                    NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std),
                    EnsureTyped(keys=["label", "psa"], dtype=torch.float32),
                    Transposed(keys=["image"], indices=(0, 3, 1, 2)),
                    DeleteItemsd(keys=[ "dwi", "adc", "mask"]),
                    ToTensord(keys=["image", "label", "smooth_mask", "psa"]),
                ]
            )
        else:
            transform = Compose(
                [
                    LoadImaged(
                        keys=["image", "mask", "dwi", "adc","smooth_mask"],
                        reader="ITKReader",
                        ensure_channel_first=True,
                        dtype=np.float32,
                    ),
                    ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
                    ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"),
                    NormalizeIntensity_customd(keys=["image"], mask_key="mask"),
                    NormalizeIntensity_customd(keys=["dwi"], mask_key="mask"),
                    ConcatItemsd(
                        keys=["image", "dwi", "adc"], name="image", dim=0
                    ),  # stacks to (3, H, W)
                    RandCropByPosNegLabeld(
                        keys=["image", "smooth_mask"],
                        label_key="smooth_mask",
                        spatial_size=(args.tile_size, args.tile_size, args.depth),
                        pos=1,
                        neg=0,
                        num_samples=args.tile_count,
                    ),
                    EnsureTyped(keys=["label"], dtype=torch.float32),
                    Transposed(keys=["image"], indices=(0, 3, 1, 2)),
                    DeleteItemsd(keys=[ "dwi", "adc", "mask"]),
                    ToTensord(keys=["image", "label", "smooth_mask"]),
                ]
            )
    return transform
"""


def get_dataloader(
    args: argparse.Namespace, split: Literal["train", "test"]
) -> torch.utils.data.DataLoader:
    data_list = load_decathlon_datalist(
        data_list_file_path=args.dataset_json,
        data_list_key=split,
        base_dir=args.data_root,
    )
    data_list_updated = [{**i, "psa": i.get("psa", [0, 0])} for i in data_list]
    cache_dir_ = os.path.join(args.logdir, "cache")
    os.makedirs(os.path.join(cache_dir_, split), exist_ok=True)
    transform = data_transform(args, split)
    dataset = PersistentDataset(
        data=data_list_updated, transform=transform, cache_dir=os.path.join(cache_dir_, split)
    )
    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=(split == "train"),
        num_workers=args.workers if split == "train" else 2,
        pin_memory=True,
        multiprocessing_context="fork" if args.workers > 0 else None,
        sampler=None,
        collate_fn=list_data_collate,
    )
    return loader