"""Utilities for unit tests."""

import datetime
import os
import re
import shutil
import subprocess
import tempfile
import unittest
import uuid
from pathlib import Path
from typing import Callable

import natsort
import numpy as np
from pydicom.dataset import FileDataset, FileMetaDataset

from dosma.cli import SUPPORTED_SCAN_TYPES, parse_args
from dosma.core.fitting import monoexponential
from dosma.core.io.format_io import ImageDataFormat
from dosma.core.med_volume import MedicalVolume
from dosma.utils import env
from dosma.utils.cmd_line_utils import ActionWrapper

UNITTEST_DATA_PATH = os.environ.get(
    "DOSMA_UNITTEST_DATA_PATH", os.path.join(os.path.dirname(__file__), "../unittest-data/")
)
UNITTEST_SCANDATA_PATH = os.path.join(UNITTEST_DATA_PATH, "scans")
TEMP_PATH = os.path.join(
    UNITTEST_SCANDATA_PATH, f"temp-{str(uuid.uuid1())}-{str(uuid.uuid4())}"
)  # should be used when for writing with assert_raises clauses

SCANS = ["qdess", "mapss", "cubequant", "cones"]
SCANS_INFO = {
    "mapss": {"expected_num_echos": 7},
    "qdess": {"expected_num_echos": 2},
    "cubequant": {"expected_num_echos": 4},
    "cones": {"expected_num_echos": 4},
}

SCAN_DIRPATHS = [os.path.join(UNITTEST_SCANDATA_PATH, x) for x in SCANS]

# Decimal precision for analysis (quantitative values, etc)
DECIMAL_PRECISION = 1  # (+/- 0.1ms)

# If elastix is available
_IS_ELASTIX_AVAILABLE = None


def is_data_available():
    disable_data = os.environ.get("DOSMA_UNITTEST_DISABLE_DATA", "").lower() == "true"
    return not disable_data and os.path.isdir(UNITTEST_DATA_PATH)


def get_scan_dirpath(scan: str):
    for ind, x in enumerate(SCANS):
        if scan == x:
            return SCAN_DIRPATHS[ind]


def get_dicoms_path(fp):
    return os.path.join(fp, "dicoms")


def get_write_path(fp, data_format: ImageDataFormat):
    return os.path.join(fp, "multi-echo-write-%s" % data_format.name)


def get_read_paths(fp, data_format: ImageDataFormat):
    """Get ground truth data (produced by imageviewer like itksnap, horos, etc)"""
    base_name = os.path.join(fp, "multi-echo-gt-%s" % data_format.name)
    files_or_dirs = os.listdir(base_name)
    fd = [x for x in files_or_dirs if re.match("e[0-9]+", x)]
    files_or_dirs = natsort.natsorted(fd)

    return [os.path.join(base_name, x) for x in files_or_dirs]


def get_data_path(fp):
    return os.path.join(fp, f"data-{str(uuid.uuid1())}")


def get_expected_data_path(fp):
    return os.path.join(fp, "expected")


def is_elastix_available():
    global _IS_ELASTIX_AVAILABLE

    if _IS_ELASTIX_AVAILABLE is None:
        disable_elastix = os.environ.get("DOSMA_UNITTEST_DISABLE_ELASTIX", None)
        if disable_elastix is None:
            _IS_ELASTIX_AVAILABLE = subprocess.run(["elastix", "--help"]).returncode == 0
        else:
            _IS_ELASTIX_AVAILABLE = disable_elastix.lower() != "true"

    return _IS_ELASTIX_AVAILABLE


def num_workers() -> int:
    return int(os.environ.get("DOSMA_NUM_WORKERS", min(8, os.cpu_count())))


def requires_packages(*packages):
    """
    Decorator for functions that should only execute when
    all packages defined by *args are supported.
    """

    def _decorator(func):
        def _wrapper(*args, **kwargs):
            if all([env.package_available(x) for x in packages]):
                func(*args, **kwargs)

        return _wrapper

    return _decorator


def generate_monoexp_data(shape=None, x=None, a=1.0, b=None):
    """Generate sample monoexponetial data.
    ``a=1.0``, ``b`` is randomly generated in interval [0.1, 1.1).

    The equation is :math:`y =  a * \\exp (b*x)`.
    """
    if b is None:
        b = np.random.rand(*shape) + 0.1
    else:
        shape = b.shape
    if x is None:
        x = np.asarray([0.5, 1.0, 2.0, 4.0])
    y = [MedicalVolume(monoexponential(t, a, b), affine=np.eye(4)) for t in x]
    return x, y, a, b


def build_dummy_headers(shape, fields=None):
    """Build dummy ``pydicom.FileDataset`` headers.

    Note these headers are not dicom compliant and should not be used to write out DICOM
    files.

    Args:
        shape (int or tuple[int]): Shape of headers array.
        fields (Dict): Fields and corresponding values to use to populate the header.

    Returns:
        ndarray: Headers
    """
    if isinstance(shape, int):
        shape = (shape,)
    num_headers = np.prod(shape)
    headers = np.asarray([_build_dummy_pydicom_header(fields) for _ in range(num_headers)])
    return headers.reshape(shape)


def _build_dummy_pydicom_header(fields=None):
    """Builds dummy pydicom-based header.

    Note these headers are not dicom compliant and should not be used to write out DICOM
    files.

    Adapted from
    https://pydicom.github.io/pydicom/dev/auto_examples/input_output/plot_write_dicom.html
    """
    suffix = ".dcm"
    filename_little_endian = tempfile.NamedTemporaryFile(suffix=suffix).name

    file_meta = FileMetaDataset()
    file_meta.MediaStorageSOPClassUID = "1.2.840.10008.5.1.4.1.1.2"
    file_meta.MediaStorageSOPInstanceUID = "1.2.3"
    file_meta.ImplementationClassUID = "1.2.3.4"

    # Create the FileDataset instance (initially no data elements, but file_meta supplied).
    ds = FileDataset(filename_little_endian, {}, file_meta=file_meta, preamble=b"\0" * 128)

    # Add the data elements -- not trying to set all required here. Check DICOM standard.
    ds.PatientName = "Test^Firstname"
    ds.PatientID = "123456"

    if fields is not None:
        for k, v in fields.items():
            setattr(ds, k, v)

    # Set the transfer syntax
    ds.is_little_endian = True
    ds.is_implicit_VR = True

    # Set creation date/time
    dt = datetime.datetime.now()
    ds.ContentDate = dt.strftime("%Y%m%d")
    timeStr = dt.strftime("%H%M%S.%f")  # long format with micro seconds
    ds.ContentTime = timeStr

    return ds


class TempPathMixin(unittest.TestCase):
    """Testing helper that creates temporary path for the class."""

    data_dirpath = None

    @classmethod
    def setUpClass(cls):
        cls.data_dirpath = Path(
            os.path.join(
                get_data_path(os.path.join(UNITTEST_SCANDATA_PATH, "temp")), f"{cls.__name__}"
            )
        )
        os.makedirs(cls.data_dirpath, exist_ok=True)

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.data_dirpath)


class ScanTest(TempPathMixin):
    from dosma.scan_sequences.scans import ScanSequence

    SCAN_TYPE = ScanSequence  # override in subclasses

    dicom_dirpath = None

    def setUp(self):
        print("Testing: ", self._testMethodName)

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        if is_data_available():
            cls.dicom_dirpath = Path(
                get_dicoms_path(os.path.join(UNITTEST_SCANDATA_PATH, cls.SCAN_TYPE.NAME))
            )

    def test_has_cmd_line_actions_attr(self):
        """
        If scan can be accessed via the command line,
        verify that the scan has a ``cmd_line_actions`` method.
        """
        # if the scan is not supported via the command line, then ignore this test
        if self.SCAN_TYPE not in SUPPORTED_SCAN_TYPES:
            return

        assert hasattr(
            self.SCAN_TYPE, "cmd_line_actions"
        ), "All scans supported by command line must have `cmd_line_actions` method"

        cmd_line_actions = self.SCAN_TYPE.cmd_line_actions()
        for func, action in cmd_line_actions:
            assert isinstance(func, Callable)
            assert isinstance(action, ActionWrapper)

            func_name = func.__name__
            cls_name = self.SCAN_TYPE.__name__
            assert action.name, f"Action for `{cls_name}.{func_name}()` must have a name"
            assert action.help, f"Action for `{cls_name}.{func_name}()` must have help menu"

        assert hasattr(
            type(self), "test_cmd_line"
        ), "All scan supported in command line must have test methods `test_cmd_line`"

    def __cmd_line_helper__(self, cmdline_str: str):
        env_args = {"--num-workers": num_workers()}
        for arg, value in env_args.items():
            if arg in cmdline_str:
                continue
            cmdline_str = f"{arg} {value} {cmdline_str}"

        cmdline_input = cmdline_str.strip().split()
        parse_args(cmdline_input)