#!/usr/bin/env python3
"""Tests for blocks."""

import shutil
import sys
from os import makedirs, remove
from os.path import isfile, join

import nibabel
import numpy as np
import pytest
from pymatreader import read_mat

from nigsp import blocks
from nigsp.io import load_nifti_get_mask
from nigsp.objects import SCGraph


# ### Unit tests
def test_nifti_to_timeseries(atlastime, atlas):
    img_in = nibabel.load(atlas)
    atlas_in = img_in.get_fdata()
    ts_out, atlas_out, img_out = blocks.nifti_to_timeseries(atlastime, atlas)
    ts_in = np.unique(atlas_in)[1:, np.newaxis]
    ts_in = np.repeat(ts_in, 2, axis=-1)

    assert (ts_in == ts_out).all()
    assert (atlas_in == atlas_out).all()
    assert (img_in.header["dim"] == img_out.header["dim"]).all()

    remove(atlastime)
    remove(atlas)


@pytest.mark.parametrize("ext", [(".nii.gz"), (".csv")])
def test_export_metrics_txt(ext, sc_mtx, atlas, sdi, testdir):
    testdir = join(testdir, "testdir")
    makedirs(testdir, exist_ok=True)
    atlas_in, mask, img = load_nifti_get_mask(atlas, ndim=3)
    mtx = read_mat(sc_mtx)
    mtx = mtx["SC_avg56"]
    sdi_in = np.genfromtxt(sdi)
    ts = np.unique(atlas_in)[1:]
    ts = np.repeat(ts[..., np.newaxis], 2, axis=-1)
    gsdi = dict.fromkeys(["hilo", "himid"], sdi_in)
    scgraph = SCGraph(mtx, ts, atlas=atlas_in, img=img, sdi=sdi_in)
    blocks.export_metric(scgraph, ext, join(testdir, "molly_"))
    scgraph = SCGraph(mtx, ts, atlas=atlas_in, img=img, gsdi=gsdi)
    blocks.export_metric(scgraph, ext, join(testdir, "cesar_"))

    assert isfile(join(testdir, f"molly_sdi{ext}"))
    assert isfile(join(testdir, f"cesar_gsdi_hilo{ext}"))
    assert isfile(join(testdir, f"cesar_gsdi_himid{ext}"))

    shutil.rmtree(testdir)
    remove(sc_mtx)
    remove(atlas)
    remove(sdi)


def test_export_metrics_nifti(sc_mtx, atlas, sdi, testdir):
    testdir = join(testdir, "testdir")
    makedirs(testdir, exist_ok=True)
    atlas_in, mask, img = load_nifti_get_mask(atlas, ndim=3)
    mtx = read_mat(sc_mtx)
    mtx = mtx["SC_avg56"]
    sdi_in = np.genfromtxt(sdi)
    ts = np.unique(atlas_in)[1:]
    ts = np.repeat(ts[..., np.newaxis], 2, axis=-1)
    scgraph = SCGraph(mtx, ts, atlas=atlas_in, sdi=sdi_in)
    blocks.export_metric(scgraph, ".nii.gz", join(testdir, "molly_"))

    assert isfile(join(testdir, "molly_sdi.tsv.gz"))

    sys.modules["nibabel"] = None
    scgraph = SCGraph(mtx, ts, atlas=atlas_in, img=img, sdi=sdi_in)
    blocks.export_metric(scgraph, ".nii.gz", join(testdir, "molly_"))
    sys.modules["nibabel"] = nibabel

    assert isfile(join(testdir, "molly_sdi.tsv.gz"))
    shutil.rmtree(testdir)
    remove(sc_mtx)
    remove(atlas)
    remove(sdi)


def test_plot_metrics(atlas, sc_mtx, sdi, testdir):
    pytest.importorskip("nilearn")
    from nilearn.plotting import find_parcellation_cut_coords

    testdir = join(testdir, "testdir")
    makedirs(testdir, exist_ok=True)
    atlas_in, _, img = load_nifti_get_mask(atlas, ndim=3)
    mtx = read_mat(sc_mtx)
    mtx = mtx["SC_avg56"]
    sdi_in = np.genfromtxt(sdi)
    ts = np.unique(atlas_in)[1:]
    ts = np.repeat(ts[..., np.newaxis], 2, axis=-1)
    gsdi = dict.fromkeys(["hilo", "himid"], sdi_in)

    scgraph = SCGraph(mtx, ts, sdi=sdi_in, img=img)
    blocks.plot_metric(scgraph, join(testdir, "maite_"), atlas=img)
    assert isfile(join(testdir, join(testdir, "maite_sdi.png")))

    atlas_in = find_parcellation_cut_coords(img)
    scgraph = SCGraph(mtx, ts, sdi=sdi_in, img=img)
    blocks.plot_metric(scgraph, join(testdir, "maite_"), atlas=atlas_in)
    assert isfile(join(testdir, join(testdir, "maite_sdi.png")))

    scgraph = SCGraph(mtx, ts, gsdi=gsdi, img=img)
    blocks.plot_metric(scgraph, join(testdir, "dimitri_"), atlas=img)
    assert isfile(join(testdir, join(testdir, "dimitri_gsdi_hilo.png")))

    shutil.rmtree(testdir)
    remove(sc_mtx)
    remove(atlas)
    remove(sdi)
