# Licensed under GPL version 3 - see LICENSE.rst
import numpy as np
from scipy.stats import kstest
from astropy.table import Table, Column
from astropy.coordinates import SkyCoord
import astropy.units as u

from transforms3d.axangles import axangle2aff

import pytest

import marxs
import marxs.optics
import marxs.optics.base
import marxs.optics.aperture
import marxs.source
from marxs.utils import generate_test_photons

@pytest.fixture(autouse=True)
def photons1000():
    '''Make a list of photons parallel to optical axis'''
    mysource = marxs.source.PointSource(coords=SkyCoord(30., 30., unit="deg"))
    np.random.seed(0)
    p = mysource.generate_photons(300 * u.ks)
    mypointing = marxs.source.FixedPointing(coords=SkyCoord(30., 30., unit='deg'))
    return mypointing(p)


def test_pos4d_no_transformation():
    '''Test coordinate transforms on initialization of optical elements.

    This test checks that the appropriate transformations are done.
    It should be extended, once derived properties (e.g. a plane constructed
    from transformed points are implemented.
    '''
    oe = marxs.optics.aperture.RectangleAperture()
    assert np.all(oe.pos4d == np.eye(4))
    assert np.all(oe.geometry['center'] == np.array([0, 0, 0, 1]))
    assert np.all(oe.geometry['e_y'] == np.array([0, 1, 0, 0]))
    assert np.all(oe.geometry['e_z'] == np.array([0, 0, 1, 0]))


def test_pos4d_rotation():
    rotation = axangle2aff(np.array([1, 1, 0]), np.deg2rad(90))
    oe = marxs.optics.aperture.RectangleAperture(orientation=rotation[:3,:3])
    assert np.all(oe.pos4d == rotation)
    assert np.all(oe.geometry['center'] == np.array([0, 0, 0, 1]))
    assert np.allclose(oe.geometry['e_y'], np.array([.5, .5, 1. / np.sqrt(2), 0]))
    assert np.allclose(oe.geometry['e_z'], np.array([1., -1, 0, 0]) / np.sqrt(2))


def test_pos4d_translation():
    oe = marxs.optics.aperture.RectangleAperture(position = np.array([-2, 3, 4.3]))
    expectedpos4d = np.eye(4)
    expectedpos4d[:3, 3] = np.array([-2, 3, 4.3])
    assert np.all(oe.pos4d == expectedpos4d)
    assert np.all(oe.geometry['center'] == np.array([-2, 3, 4.3, 1]))
    assert np.all(oe.geometry['e_y'] == np.array([0, 1, 0, 0]))
    assert np.all(oe.geometry['e_z'] == np.array([0, 0, 1, 0]))



mark = pytest.mark.parametrize

all_slits = [marxs.optics.aperture.RectangleAperture(zoom=2),
             marxs.optics.aperture.RectangleAperture(zoom=[2,2,2]),
            ]

@mark('myslit', all_slits)
def test_pos4d_transforms_slit(photons1000, myslit):
    '''Test coordinate transforms on initialization of optical elements.

    The initial 4D transforms should be done to any optical element.
    Here, I pick the entrance aperture for testing, because it places the
    positional vector of the plucker coordinates in a plane, independent
    of the initial values.
    '''

    p = myslit(photons1000)
    assert np.allclose(p['pos'][:, 0], 0)
    assert kstest((p['pos'][:, 1] + 2) / 4, "uniform")[1] > 0.01
    assert kstest((p['pos'][:, 2] + 2) / 4, "uniform")[1] > 0.01


def test_pos4d_transforms_slit_rotated(photons1000):
    '''Test coordinate transforms on rotated entrance aperture.'''
    p = photons1000

    rotation = axangle2aff(np.array([0, 1, 0]), np.deg2rad(90))
    myslit = marxs.optics.aperture.RectangleAperture(orientation=rotation[:3, :3], zoom=0.5)
    p = myslit(p)
    assert np.allclose(p['pos'][:, 2], 0)
    assert kstest(p['pos'][:, 0] + 0.5, "uniform")[1] > 0.01
    assert kstest(p['pos'][:, 1] + 0.5, "uniform")[1] > 0.01
    # Sometimes fails, because rays exactly parallel to aperture.
    # Find other example, e.g. 45 deg.
    # Check in aperture if that's the case and raise warning.


def test_photonlocalcoords_decorator():

    pos = np.array([[1, 0, 0, 1], [0, 1, 0, 1]])
    dir = np.array([[1, 0, 0, 0], [0, 1, 0, 0]])
    photontab = Table([Column(data=pos, name='pos'),
                       Column(data=dir, name='dir')
                       ])

    class OEforTest(marxs.optics.base.OpticalElement):
        @marxs.optics.base.photonlocalcoords
        def functiontodecorate(self, photons):
            assert np.allclose(photons['dir'], dir)
            assert np.allclose(photons['pos'], np.array([[2., 2, 3, 1], [1, 3, 3, 1]]))
            return photons

    oetest = OEforTest(position=np.array([-1., -2, -3]))
    photontab = oetest.functiontodecorate(photontab)
    assert np.allclose(photontab['dir'], dir)
    assert np.allclose(photontab['pos'], pos)


def test_FlatStack():
    '''Run a stack of two elements and check that both are applied to the photons.'''
    fs = marxs.optics.FlatStack(position=[0, 5, 0], zoom=2,
                                elements=[marxs.optics.EnergyFilter, marxs.optics.FlatDetector],
                                keywords=[{'filterfunc': lambda x: 0.5},{}])
    # Check all layers are at the same position
    assert np.allclose(fs.geometry['center'], np.array([0, 5, 0, 1]))
    assert np.allclose(fs.elements[0].geometry['center'], np.array([0, 5, 0, 1]))
    assert np.allclose(fs.elements[1].geometry['center'], np.array([0, 5, 0, 1]))
    assert np.allclose(fs.pos4d, fs.elements[1].pos4d)

    p = generate_test_photons(5)
    # Photons 0, 1 miss the stack
    p['pos'][:, 1] = [0, 1, 3.1, 4, 5]
    fs.loc_coos_name = ['a', 'b']
    p = fs(p)
    assert np.allclose(p['probability'], [1, 1, .5, .5, .5])
    assert np.all(np.isnan(p['a'][:2]))
    assert np.allclose(p['a'][2:], [-1.9, -1., 0])
