"""
OXASL - Module to generate additional ROIs for analysis

Copyright (c) 2008 University of Nottingham
"""
import numpy as np

from fsl.data.image import Image
from fsl.data.atlases import AtlasRegistry

from oxasl import reg

def run(wsp):
    """
    Generate additional ROIs used for reporting and region analysis

    wsp.rois.gm_asl: GM mask in ASL space (>0.5 probability)
    wsp.rois.wm_asl: WM mask in ASL space (>0.5 probability)
    wsp.rois.pure_gm_asl: 'pure' GM mask in ASL space (default > 0.8 probability, but uses pure_gm_thresh option)
    wsp.rois.pure_wm_asl: 'pure' WM mask in ASL space (default > 0.9 probability, but uses pure_gm_thresh option)
    wsp.rois.cortial_gm_asl: Cortical GM mask in ASL space, same as pure_gm_asl but masked to remove subcortical structures
    wsp.rois.cerebral_wm_asl: Cererbral WM mask in ASL space, same as pure_wm_asl but masked to remove subcortical structures
    """
    if wsp.structural.struc is not None:
        # Get the cortex 
        atlases = AtlasRegistry()
        atlases.rescanAtlases()
        atlas = atlases.loadAtlas("harvardoxford-subcortical", loadSummary=False, resolution=2)
        wsp.rois.cortex_std = Image(np.mean(atlas.data[..., 0:2] + atlas.data[..., 11:13], axis=-1) * 2, header=atlas.header)
        cortex_asl = reg.change_space(wsp, wsp.rois.cortex_std, "asl")
        wsp.rois.cortex_asl = Image((cortex_asl.data > 50).astype(int), header=cortex_asl.header)

        if wsp.structural.gm_pv_asl is None:
            wsp.structural.gm_pv_asl = reg.change_space(wsp, wsp.structural.gm_pv, "asl")
            wsp.structural.wm_pv_asl = reg.change_space(wsp, wsp.structural.wm_pv, "asl")
            wsp.structural.csf_pv_asl = reg.change_space(wsp, wsp.structural.csf_pv, "asl")

        gm = np.asarray(wsp.structural.gm_pv_asl.data)
        wm = np.asarray(wsp.structural.wm_pv_asl.data)

        wsp.rois.pure_gm_thresh = wsp.ifnone("pure_gm_thresh", 0.8)
        wsp.rois.pure_wm_thresh = wsp.ifnone("pure_wm_thresh", 0.9)
        some_gm, some_wm = gm > 0.5, wm > 0.5
        pure_gm, pure_wm = gm > wsp.rois.pure_gm_thresh, wm > wsp.rois.pure_wm_thresh
        wsp.rois.gm_asl = Image(some_gm.astype(int), header=wsp.structural.gm_pv_asl.header)
        wsp.rois.pure_gm_asl = Image(pure_gm.astype(int), header=wsp.structural.gm_pv_asl.header)
        wsp.rois.wm_asl = Image(some_wm.astype(int), header=wsp.structural.wm_pv_asl.header)
        wsp.rois.pure_wm_asl = Image(pure_wm.astype(int), header=wsp.structural.wm_pv_asl.header)
        wsp.rois.cortical_gm_asl = Image(np.logical_and(pure_gm, cortex_asl.data).astype(int), header=wsp.structural.gm_pv_asl.header)
        wsp.rois.cerebral_wm_asl = Image(np.logical_and(pure_wm, cortex_asl.data).astype(int), header=wsp.structural.wm_pv_asl.header)
