#!/usr/bin/env fslpython

# test avscale with different input affines


import sys
import os

import fsl.utils.run as run

import numpy as np


PI    = np.pi
PION2 = np.pi / 2


# Each test has the form (input-affine, expected-output)
# where expected-output comprises:
#   - scales
#   - translations
#   - rotations
#   - skews
#   - L/R orientation (preserved or swapped)
tests = [

    ([[1, 0, 0, 0],
      [0, 1, 0, 0],
      [0, 0, 1, 0],
      [0, 0, 0, 1]], [(1, 1, 1), (0, 0, 0), (0, 0, 0), (0, 0, 0), 'preserved']),
    ([[2, 0, 0, 0],
      [0, 2, 0, 0],
      [0, 0, 2, 0],
      [0, 0, 0, 1]], [(2, 2, 2), (0, 0, 0), (0, 0, 0), (0, 0, 0), 'preserved']),
    ([[-2, 0, 0, 0],
      [ 0, 2, 0, 0],
      [ 0, 0, 2, 0],
      [ 0, 0, 0, 1]], [(2, 2, 2), (0, 0, 0), (0, 0, PI), (0, 0, 0), 'swapped']),
    ([[0, 0, 1, 0],
      [0, 1, 0, 0],
      [1, 0, 0, 0],
      [0, 0, 0, 1]], [(1, 1, 1), (0, 0, 0), (0, -PION2, 0), (0, 0, 0), 'swapped']),
    ([[ 0,  0, -1, 0],
      [ 0, -1,  0, 0],
      [-1,  0,  0, 0],
      [ 0,  0,  0, 1]], [(1, 1, 1), (0, 0, 0), (-PI, PION2, 0), (0, 0, 0), 'preserved']),
    ([[0, 1, 0, 0],
      [1, 0, 0, 0],
      [0, 0, 1, 0],
      [0, 0, 0, 1]], [(1, 1, 1), (0, 0, 0), (0, 0, PION2), (0, 0, 0), 'swapped']),
    ([[1, 0, 0, 0],
      [0, 0, 1, 0],
      [0, 1, 0, 0],
      [0, 0, 0, 1]], [(1, 1, 1), (0, 0, 0), (PION2, 0, 0), (0, 0, 0), 'swapped']),
    ([[-2, 0, 0,  90],
      [ 0, 2, 0, -126],
      [ 0, 0, 2, -72],
      [ 0, 0, 0,  1]], [(2, 2, 2), (90, -126, -72), (0, 0, PI), (0, 0, 0), 'swapped']),

]

def read_avscale_output(output):
    lines = output.split('\n')
    # scales
    # translations
    # rotations
    # skews
    # l/r orientation (preserved or swapped)
    scales       = [l for l in lines if l.startswith('Scales ')]        [0]
    translations = [l for l in lines if l.startswith('Translations ')]  [0]
    rotations    = [l for l in lines if l.startswith('Rotation Angles')][0]
    skews        = [l for l in lines if l.startswith('Skews ')]         [0]
    orient       = [l for l in lines if l.startswith('Left-Right ')]    [0]

    scales       = [float(s) for s in scales      .split()[-3:]]
    translations = [float(s) for s in translations.split()[-3:]]
    rotations    = [float(s) for s in rotations   .split()[-3:]]
    skews        = [float(s) for s in skews       .split()[-3:]]
    orient       = orient.split()[-1]

    return {
        'scales'       : np.array(scales),
        'translations' : np.array(translations),
        'rotations'    : np.array(rotations),
        'skews'        : np.array(skews),
        'orient'       : orient.lower(),
    }


def test_avscale():

    for affine, expected in tests:
        affine = np.array(affine)
        np.savetxt('affine.mat', affine, '%0.6f')
        print(affine)
        output = run.runfsl('avscale --allparams affine.mat')
        output = read_avscale_output(output)

        scales, translations, rotations, skews, orient = expected

        try:
            assert np.all(np.isclose(scales,       output['scales']))
            assert np.all(np.isclose(translations, output['translations']))
            assert np.all(np.isclose(rotations,    output['rotations']))
            assert np.all(np.isclose(skews,        output['skews']))
            assert                   orient     == output['orient']
        except AssertionError:
            print(affine)
            print(output)
            raise


if __name__ == '__main__':
    if len(sys.argv) > 1:
        os.chdir(sys.argv[1])
    test_avscale()