"""
Compute reconstructions with CTAC and DLAC (attenuation maps created by `compute_mu_maps`).

This script implements the (post-reconstructions attenuation correction) PRAC algorithm
by first projection a non-attenuation-corrected reconstruction and then reconstructing this
projection with an attenuation map.
"""
import argparse
import os

import pydicom
import stir

stir.Verbosity_set(0)

from mu_map.file.dicom import change_uid, load_dcm, update_dcm
from mu_map.file.util import load_as_interfile
from mu_map.logging import add_logging_args, get_logger_by_args
from mu_map.recon.filter import GaussianFilter
from mu_map.recon.osem import reconstruct
from mu_map.recon.project import forward_project

parser = argparse.ArgumentParser()
parser.add_argument(
    "--mu_maps",
    type=str,
    required=True,
    help="directory of mu maps computed with `mu_map.scripts.compute_mu_maps`",
)
parser.add_argument(
    "--out_dir",
    type=str,
    required=True,
    help="directory where resulting reconstructions are written to",
)
parser.add_argument(
    "--description_postfix",
    type=str,
    help="postfix for the DICOM SeriesDescription",
)
parser.add_argument(
    "--dataset_dir",
    type=str,
    default="data/second/",
    help="directory of the dataset with reconstructions and real mu maps",
)
parser.add_argument(
    "--images_dir",
    type=str,
    default="images",
    help="directory under `--dataset_dir` actually containing the images",
)
parser.add_argument(
    "--n_subsets",
    type=int,
    default=3,
    help="number of subsets used for the OSEM reconstruction",
)
parser.add_argument(
    "--n_iterations",
    type=int,
    default=10,
    help="number of iterations performed by the OSEM reconstruction",
)
add_logging_args(parser, defaults={"--logfile": "compute_recons.log"})
args = parser.parse_args()

args.images_dir = os.path.join(args.dataset_dir, args.images_dir)
args.description_postfix = (
    "" if args.description_postfix is None else f" - {args.description_postfix}"
)
args.logfile = os.path.join(args.out_dir, args.logfile)

logger = get_logger_by_args(args)
logger.info(args)

files_mu_maps_dl = os.listdir(args.mu_maps)
files_mu_maps_dl = filter(lambda f: f.split(".")[-1] == "dcm", files_mu_maps_dl)
files_mu_maps_dl = list(files_mu_maps_dl)
files_mu_maps_dl.sort()

for file_mu_map_dl in files_mu_maps_dl:
    logger.info(f"Process {file_mu_map_dl}")

    file_recon_ac = file_mu_map_dl.replace("mu_map_dl", "recon_ac_nsc")
    file_recon_ac = os.path.join(args.images_dir, file_recon_ac)
    recon_ac_dcm, _ = load_dcm(file_recon_ac, direction=1)
    logger.info(f"Found AC recon: {file_recon_ac}")

    file_recon_nac = file_mu_map_dl.replace("mu_map_dl", "recon_nac_nsc")
    file_recon_nac = os.path.join(args.images_dir, file_recon_nac)
    recon_nac = load_as_interfile(file_recon_nac, direction=1)
    logger.info(f"Found nAC recon: {file_recon_nac}")

    file_mu_map_ct = file_mu_map_dl.replace("mu_map_dl", "mu_map")
    file_mu_map_ct = os.path.join(args.images_dir, file_mu_map_ct)
    mu_map_ct = load_as_interfile(file_mu_map_ct, direction=1)
    logger.info(f"Found CT-based mu map: {file_mu_map_ct}")

    mu_map_dl = load_as_interfile(
        os.path.join(args.mu_maps, file_mu_map_dl), direction=1
    )

    for mu_map, label in zip([mu_map_ct, mu_map_dl], ["CTAC", "DLAC"]):
        logger.info(f"Compute {label} reconstruction")
        projection = forward_project(recon_nac, n_slices=mu_map.image.shape[0])
        _recon_ac = reconstruct(
            projection,
            mu_map=mu_map,
            n_subsets=args.n_subsets,
            n_iterations=args.n_iterations,
            postfilter=GaussianFilter(projection),
        )
        # for some reason STIR reconstruction are inverted in the x-direction
        # this is reverted here
        _recon_ac.image = _recon_ac.image[:, :, ::-1]

        _recon_ac_dcm = change_uid(update_dcm(recon_ac_dcm.copy(), _recon_ac.image))
        _recon_ac_dcm.SeriesDescription = _recon_ac_dcm.SeriesDescription.replace(
            "NoSC - AC ", f"{label}{args.description_postfix}"
        )
        _file_recon_ac = file_mu_map_dl.replace("mu_map_dl", f"recon_{label.lower()}")
        _file_recon_ac = os.path.join(args.out_dir, _file_recon_ac)
        pydicom.dcmwrite(_file_recon_ac, _recon_ac_dcm)
        logger.info(f"Store {label} recon at {_file_recon_ac}")