import argparse
import logging
import os
import shutil

import nrrd
import numpy as np
import SimpleITK as sitk
from tqdm import tqdm


def smoothen_mask(args, sigma=1.0):
    files = os.listdir(args.seg_dir)

    for file in files:
        mask = sitk.ReadImage(os.path.join(args.seg_dir, file))
        mask = sitk.Cast(mask, sitk.sitkFloat32)

        # Smooth
        smoothed = sitk.DiscreteGaussian(mask, variance=1.5)

        # Threshold back to binary
        smooth_mask = smoothed > 0.5
        smooth_mask = sitk.Cast(smooth_mask, sitk.sitkUInt8)

        out_path = os.path.join(args.smooth_seg_dir_temp, file)
        sitk.WriteImage(smooth_mask, out_path)


def get_heatmap(args: argparse.Namespace) -> argparse.Namespace:
    """
    Generate heatmaps from DWI (Diffusion Weighted Imaging) and ADC (Apparent Diffusion Coefficient) medical imaging data.
    This function processes medical imaging files (DWI and ADC) along with their corresponding
    segmentation masks to create normalized heatmaps. It combines the DWI
    and ADC heatmaps through element-wise multiplication.
    Args:
        args: An object containing the following attributes:
            - t2_dir (str): Directory path containing T2 image files.
            - dwi_dir (str): Directory path containing DWI image files.
            - adc_dir (str): Directory path containing ADC image files.
            - seg_dir (str): Directory path containing segmentation mask files.
            - output_dir (str): Base output directory where 'heatmaps/' subdirectory will be created.
            - heatmapdir (str): Output directory for heatmap files (created by function).
    Returns:
        args: The modified args object with heatmapdir attribute set.
    Raises:
        FileNotFoundError: If input directories or files do not exist.
        ValueError: If NRRD files cannot be read properly.
    Notes:
        - DWI heatmap is normalized as (dwi - min) / (max - min)
        - ADC heatmap is normalized as (max - adc) / (max - min) (inverted)
        - Final heatmap is re-normalized to [0, 1] range
        - If all values in a mask region are identical, the heatmap is skipped for that modality
        - Output files are written in NRRD format with the same header as the input DWI file
    """

    files = os.listdir(args.t2_dir)
    args.heatmapdir = os.path.join(args.output_dir, "heatmaps/")
    os.makedirs(args.heatmapdir, exist_ok=True)
    args.smooth_seg_dir_temp = os.path.join(args.output_dir, "smooth_prostate_mask_temp/")
    os.makedirs(args.smooth_seg_dir_temp, exist_ok=True)
    args.smooth_seg_dir = os.path.join(args.output_dir, "smooth_prostate_mask/")
    os.makedirs(args.smooth_seg_dir, exist_ok=True)
    smoothen_mask(args)

    logging.info("Starting heatmap generation")
    for file in tqdm(files):
        bool_dwi = False
        bool_adc = False
        mask_temp, header_mask = nrrd.read(os.path.join(args.seg_dir, file))
        # spacing = np.linalg.norm(header_mask['space directions'], axis=1)
        dwi, header_dwi = nrrd.read(os.path.join(args.dwi_dir, file))
        adc, header_adc = nrrd.read(os.path.join(args.adc_dir, file))
        nonzero_vals_dwi = dwi[mask_temp > 0]
        mask, _ = nrrd.read(os.path.join(args.smooth_seg_dir_temp, file))
        mask = np.maximum(mask, mask_temp)

        if len(nonzero_vals_dwi) > 0:
            # min_val = nonzero_vals_dwi.min()
            # max_val = nonzero_vals_dwi.max()
            min_val, max_val = np.percentile(nonzero_vals_dwi, [1, 99])
            clipped_dwi = np.clip(dwi, min_val, max_val)

            heatmap_dwi = np.zeros_like(clipped_dwi, dtype=np.float32)

            if min_val != max_val:
                heatmap_dwi = (clipped_dwi - min_val) / (max_val - min_val)
                masked_heatmap_dwi = np.where(mask > 0, heatmap_dwi, 0)
            else:
                bool_dwi = True

        else:
            bool_dwi = True

        nonzero_vals_adc = adc[mask > 0]

        if len(nonzero_vals_adc) > 0:
            min_val = 0
            max_val = 1
            heatmap_adc = np.zeros_like(adc, dtype=np.float32)
            heatmap_adc = (max_val - adc) / (max_val - min_val)
            masked_heatmap_adc = np.where(mask > 0, heatmap_adc, 0)

        else:
            bool_adc = True

        if not bool_dwi and not bool_adc:
            mix_mask = (masked_heatmap_dwi * 0.3) + (masked_heatmap_adc * 0.7)
            # mix_mask = (masked_heatmap_dwi**0.5) * (masked_heatmap_adc**2.0)
            write_header = header_dwi
        elif bool_dwi:
            mix_mask = masked_heatmap_adc
            write_header = header_adc
        else:
            mix_mask = np.ones_like(adc, dtype=np.float32)
            write_header = header_dwi

        mix_mask = (mix_mask - mix_mask[mask > 0].min()) / (
            mix_mask[mask > 0].max() - mix_mask[mask > 0].min()
        )
        mix_mask = np.where(mask > 0, mix_mask, 0)
        nrrd.write(os.path.join(args.heatmapdir, file), mix_mask, write_header)
        nrrd.write(os.path.join(args.smooth_seg_dir, file), mask, header_mask)

    shutil.rmtree(args.smooth_seg_dir_temp)

    return args