import argparse
import logging
import os

import nrrd
import numpy as np
from skimage import exposure
from tqdm import tqdm


def get_histmatched(
    data: np.ndarray, ref_data: np.ndarray, mask: np.ndarray, ref_mask: np.ndarray
) -> np.ndarray:
    """
    Perform histogram matching on source data using a reference image.
    This function adjusts the histogram of the source image to match the
    histogram of the reference image within masked regions. Only pixels
    where the mask is greater than 0 are considered for matching.
    Args:
        data: Source image array to be histogram matched.
        ref_data: Reference image array whose histogram will be used as target.
        mask: Binary mask for source image indicating valid pixels (values > 0).
        ref_mask: Binary mask for reference image indicating valid pixels (values > 0).
    Returns:
        Histogram-matched image with the same shape as input data.
        Only pixels in masked regions are modified; unmasked pixels remain unchanged.
    Example:
        >>> matched = get_histmatched(source_img, reference_img, source_mask, ref_mask)
    """
    source_pixels = data[mask > 0]
    ref_pixels = ref_data[ref_mask > 0]
    matched_pixels = exposure.match_histograms(source_pixels, ref_pixels)
    matched_img = data.copy()
    matched_img[mask > 0] = matched_pixels

    return matched_img


def histmatch(args: argparse.Namespace) -> argparse.Namespace:
    files = os.listdir(args.t2_dir)

    t2_histmatched_dir = os.path.join(args.output_dir, "t2_histmatched")
    dwi_histmatched_dir = os.path.join(args.output_dir, "DWI_histmatched")
    adc_histmatched_dir = os.path.join(args.output_dir, "ADC_histmatched")
    os.makedirs(t2_histmatched_dir, exist_ok=True)
    os.makedirs(dwi_histmatched_dir, exist_ok=True)
    os.makedirs(adc_histmatched_dir, exist_ok=True)
    logging.info("Starting histogram matching")

    for file in tqdm(files):
        t2_image, header_t2 = nrrd.read(os.path.join(args.t2_dir, file))
        dwi_image, header_dwi = nrrd.read(os.path.join(args.dwi_dir, file))
        adc_image, header_adc = nrrd.read(os.path.join(args.adc_dir, file))

        ref_t2, _ = nrrd.read(os.path.join(args.project_dir, "dataset", "t2_reference.nrrd"))
        ref_dwi, _ = nrrd.read(os.path.join(args.project_dir, "dataset", "dwi_reference.nrrd"))
        ref_adc, _ = nrrd.read(os.path.join(args.project_dir, "dataset", "adc_reference.nrrd"))
        prostate_mask, _ = nrrd.read(os.path.join(args.seg_dir, file))
        ref_prostate_mask, _ = nrrd.read(
            os.path.join(args.project_dir, "dataset", "prostate_segmentation_reference.nrrd")
        )

        histmatched_t2 = get_histmatched(t2_image, ref_t2, prostate_mask, ref_prostate_mask)
        histmatched_dwi = get_histmatched(dwi_image, ref_dwi, prostate_mask, ref_prostate_mask)
        histmatched_adc = get_histmatched(adc_image, ref_adc, prostate_mask, ref_prostate_mask)

        nrrd.write(os.path.join(t2_histmatched_dir, file), histmatched_t2, header_t2)
        nrrd.write(os.path.join(dwi_histmatched_dir, file), histmatched_dwi, header_dwi)
        nrrd.write(os.path.join(adc_histmatched_dir, file), histmatched_adc, header_adc)

        args.t2_dir = t2_histmatched_dir
        args.dwi_dir = dwi_histmatched_dir
        args.adc_dir = adc_histmatched_dir

    return args