import numpy

from ophyd import PseudoSingle, SoftPositioner, Component as Cpt

from addams_bec.devices.diffract import X04V, X04H

class SimulatedX04V(X04V):

    ov = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0., kind="normal")
    alp = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0., kind="normal")
    delta = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0., kind="normal")
    gam = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0., kind="normal")
    nu = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0, kind="normal")


class SimulatedX04H(X04H):

    phi = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0., kind="normal")
    oh = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0., kind="normal")
    delta = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0., kind="normal")
    gam = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0., kind="normal")
    nu = Cpt(SoftPositioner, limits=(-180, 180), init_pos=0., kind="normal")


def test_UB_x04v():
    x04v = SimulatedX04V(prefix='', name="x04v")

    # Compute UB matrix
    x04v.calc.lattice = (5.43, 5.5, 7.71, 90, 90, 90)
    x04v.calc.wavelength = 0.71
    x04v.calc.add_reflection(1, -1, 2, numpy.array([ 88.2400, 10,  8.7000, 12.3200]))

    # Compute UB matrix
    x04v.calc.lattice = (5.43, 5.5, 7.71, 90, 90, 90)
    x04v.calc.wavelength = 0.71
    x04v.calc.add_reflection(1, -1, 2, numpy.array([ 88.2400, 10,  8.7000, 12.3200]))
    x04v.calc.add_reflection(2, -2, 0, numpy.array([ 35.3500, 10, 17.3502, 12.4000]))
    x04v.calc.add_reflection(2, -2, 4, numpy.array([101.4501, 10, 17.1200, 25.3300]))
    x04v.calc.add_reflection(6, -2, 4, numpy.array([ 95.6500, 10, 46.1300, 31.6600]))
    x04v.calc.add_reflection(1, -3, 1, numpy.array([ 16.2754, 10, 11.2367, 21.6627]))
    x04v.calc.compute_UB()

    UBspec = numpy.array([
        [0.58641,   -0.56745,   -0.56980],
        [0.99701,    0.35123,    0.34255],
        [0.00667,   -0.93809,    0.47113],
    ])

    assert(numpy.allclose(UBspec, x04v.calc.UB, atol=1e-4))

    x04v.move(h=1)

def test_UB_x04h():
    x04h = SimulatedX04H(prefix='', name='x04h')
    x04h.calc.lattice = (2.989, 2.989, 2.989, 90, 90, 90)
    x04h.calc.wavelength = 1.631

    x04h.calc.add_reflection(1,  1,  1, numpy.array([ -1.0000,    5.0000,   30.5758,   50.0038]))
    x04h.calc.add_reflection(1, -1, -1, numpy.array([ 89.2000,    5.0000,   30.5758,   50.0038]))
    x04h.calc.add_reflection(1, -1,  1, numpy.array([179.1800,    5.0000,   30.5758,   50.0038]))
    x04h.calc.compute_UB()

    UBspec = numpy.array([
        [-2.04607,  -2.51605,    2.04855],
        [ 0.46587,  -1.58340,   -0.46982],
        [ 2.10211,   0.00000,    0.00000],
    ])

    assert(numpy.allclose(UBspec, x04h.calc.UB, atol=1e-4))

    x04h.move(h=1)