import os
import logging
import nibabel as nib
import numpy as np
from scipy.ndimage import binary_dilation
from memori.pathman import append_suffix, repath
from memori.helpers import create_output_path
from omni.interfaces.common import normalize
from omni.interfaces.afni import Autobox, resample
from omni.preprocessing import saturate, localized_contrast_enhance


@create_output_path
def weight_mask_and_autobox(
    output_path: str, t1_debias: str, t2_debias: str, anat_bet_mask: str, anat_eye_mask: str, dilation_size: int = 30
):
    """Generate an anatomical weight mask

    Parameters
    ----------
    output_path : str
        Output path to write out files to
    t1_debias : str
        T1 image.
    t2_debias : str
        T2 image.
    anat_bet_mask : str
        Anatomical brain mask
    anat_eye_mask : str
        Anatomical eye mask
    dilation_size : int
        Dilation size for the weight mask

    Returns
    -------
    str
        Weight mask
    str
        Weight mask + autobox
    str
        Brain mask + autobox
    str
        Debiased T1 + autobox
    str
        Debiased T2 + autobox
    str
        Debiased T1 + autobox + sat/lce
    str
        Debiased T2 + autobox + sat/lce
    str
        Debiased T1 + autobox + sat/lce + bet
    str
        Debiased T2 + autobox + sat/lce + bet
    """
    logging.info("Creating anatomical weight mask...")
    logging.info(
        """
        It may look like that this part has hung,
        but it really is just that slow... Blame scipy for having a
        crappy mathematical morphology library..."""
    )

    # load bet masks
    bet_mask_img = nib.load(anat_bet_mask)
    bet_data = bet_mask_img.get_fdata()
    eye_mask_img = nib.load(anat_eye_mask)
    eye_data = eye_mask_img.get_fdata()

    # create structural element for dilation
    strel = np.ones(tuple([dilation_size for _ in range(3)]), dtype=np.bool)

    # create weight mask by adding dilated mask + eye_masl
    weight_mask_data = normalize(
        np.logical_or(binary_dilation(bet_data.astype("?"), strel), eye_data.astype("?")).astype(bet_data.dtype)
        + bet_data
    )

    # save the mask to file
    weight_mask = os.path.join(output_path, "weight_mask.nii.gz")
    nib.Nifti1Image(weight_mask_data, bet_mask_img.affine, bet_mask_img.header).to_filename(weight_mask)

    # autobox the weighted anatomical mask
    logging.info("Creating autobox mask...")
    weight_mask_ab = os.path.join(output_path, "weight_mask_ab.nii.gz")
    Autobox(weight_mask_ab, weight_mask)

    # resample each t1/t2 images to autobox grid
    logging.info("Autoboxing images...")
    t1_debias_ab = repath(output_path, append_suffix(t1_debias, "_ab"))
    t2_debias_ab = repath(output_path, append_suffix(t2_debias, "_ab"))
    anat_bet_mask_ab = repath(output_path, append_suffix(anat_bet_mask, "_ab"))
    resample(t1_debias_ab, weight_mask_ab, t1_debias)
    resample(t2_debias_ab, weight_mask_ab, t2_debias)
    resample(anat_bet_mask_ab, weight_mask_ab, anat_bet_mask)

    # saturate and lce each image
    logging.info("Saturate and localize contrast enhance images...")
    t1_debias_ab_sat_lce = append_suffix(t1_debias_ab, "_sat_lce")
    t2_debias_ab_sat_lce = append_suffix(t2_debias_ab, "_sat_lce")
    bet_mask_ab_img = nib.load(anat_bet_mask_ab)
    t1_debias_ab_img = nib.load(t1_debias_ab)
    t2_debias_ab_img = nib.load(t2_debias_ab)
    t1_debias_ab_sat_img = saturate(t1_debias_ab_img)
    t2_debias_ab_sat_img = saturate(t2_debias_ab_img)
    t1_debias_ab_sat_lce_img = localized_contrast_enhance(t1_debias_ab_sat_img, bet_mask_ab_img)
    t2_debias_ab_sat_lce_img = localized_contrast_enhance(t2_debias_ab_sat_img, bet_mask_ab_img)
    t1_debias_ab_sat_lce_img.to_filename(t1_debias_ab_sat_lce)
    t2_debias_ab_sat_lce_img.to_filename(t2_debias_ab_sat_lce)

    # get bet versions of sat/lce images
    t1_debias_ab_sat_lce_bet = append_suffix(t1_debias_ab_sat_lce, "_bet")
    t2_debias_ab_sat_lce_bet = append_suffix(t2_debias_ab_sat_lce, "_bet")
    nib.Nifti1Image(
        t1_debias_ab_sat_lce_img.get_fdata() * bet_mask_ab_img.get_fdata(),
        t1_debias_ab_sat_lce_img.affine,
        t1_debias_ab_sat_lce_img.header,
    ).to_filename(t1_debias_ab_sat_lce_bet)
    nib.Nifti1Image(
        t2_debias_ab_sat_lce_img.get_fdata() * bet_mask_ab_img.get_fdata(),
        t2_debias_ab_sat_lce_img.affine,
        t2_debias_ab_sat_lce_img.header,
    ).to_filename(t2_debias_ab_sat_lce_bet)

    # return masks and images
    return (
        weight_mask,
        weight_mask_ab,
        anat_bet_mask_ab,
        t1_debias_ab,
        t2_debias_ab,
        t1_debias_ab_sat_lce,
        t2_debias_ab_sat_lce,
        t1_debias_ab_sat_lce_bet,
        t2_debias_ab_sat_lce_bet,
    )
