import time
import logging

import pytest
import operator
import numpy as np


from copy import copy

from ophyd.epics_motor import EpicsMotor
from ophyd.pseudopos import (PseudoPositioner, PseudoSingle)

from ophyd import (Component as C)
from ophyd.utils import ExceptionBundle

logger = logging.getLogger(__name__)


def setUpModule():
    logging.getLogger('ophyd.pseudopos').setLevel(logging.DEBUG)


def tearDownModule():

    logger.debug('Cleaning up')
    logging.getLogger('ophyd.pseudopos').setLevel(logging.INFO)


motor_recs = ['XF:31IDA-OP{Tbl-Ax:X1}Mtr',
              'XF:31IDA-OP{Tbl-Ax:X2}Mtr',
              'XF:31IDA-OP{Tbl-Ax:X3}Mtr',
              'XF:31IDA-OP{Tbl-Ax:X4}Mtr',
              'XF:31IDA-OP{Tbl-Ax:X5}Mtr',
              'XF:31IDA-OP{Tbl-Ax:X6}Mtr',
              ]


class Pseudo3x3(PseudoPositioner):
    pseudo1 = C(PseudoSingle, '', limits=(-10, 10), egu='a')
    pseudo2 = C(PseudoSingle, '', limits=(-10, 10), egu='b')
    pseudo3 = C(PseudoSingle, '', limits=None, egu='c')
    real1 = C(EpicsMotor, motor_recs[0])
    real2 = C(EpicsMotor, motor_recs[1])
    real3 = C(EpicsMotor, motor_recs[2])

    def forward(self, pseudo_pos):
        pseudo_pos = self.PseudoPosition(*pseudo_pos)
        # logger.debug('forward %s', pseudo_pos)
        return self.RealPosition(real1=-pseudo_pos.pseudo1,
                                 real2=-pseudo_pos.pseudo2,
                                 real3=-pseudo_pos.pseudo3)

    def inverse(self, real_pos):
        real_pos = self.RealPosition(*real_pos)
        # logger.debug('inverse %s', real_pos)
        return self.PseudoPosition(pseudo1=real_pos.real1,
                                   pseudo2=real_pos.real2,
                                   pseudo3=real_pos.real3)


class Pseudo1x3(PseudoPositioner):
    pseudo1 = C(PseudoSingle, limits=(-10, 10))
    real1 = C(EpicsMotor, motor_recs[0])
    real2 = C(EpicsMotor, motor_recs[1])
    real3 = C(EpicsMotor, motor_recs[2])

    def forward(self, pseudo_pos):
        pseudo_pos = self.PseudoPosition(*pseudo_pos)
        # logger.debug('forward %s', pseudo_pos)
        return self.RealPosition(real1=-pseudo_pos.pseudo1,
                                 real2=-pseudo_pos.pseudo1,
                                 real3=-pseudo_pos.pseudo1)

    def inverse(self, real_pos):
        real_pos = self.RealPosition(*real_pos)
        # logger.debug('inverse %s', real_pos)
        return self.PseudoPosition(pseudo1=-real_pos.real1)


class FaultyStopperEpicsMotor(EpicsMotor):
    def stop(self, *, success=False):
        raise RuntimeError('Expected exception')


class FaultyPseudo1x3(Pseudo1x3):
    real1 = C(FaultyStopperEpicsMotor, motor_recs[0])


def test_onlypseudo():
    # can't instantiate it on its own
    with pytest.raises(TypeError):
        PseudoPositioner('prefix')


def test_position_wrapper():
    pseudo = Pseudo3x3('', name='mypseudo', concurrent=False)

    test_pos = pseudo.PseudoPosition(pseudo1=1, pseudo2=2, pseudo3=3)
    extra_kw = dict(a=3, b=4, c=6)

    # positional arguments
    assert pseudo.to_pseudo_tuple(1, 2, 3, **extra_kw) == (test_pos, extra_kw)
    # sequence
    assert pseudo.to_pseudo_tuple((1, 2, 3), **extra_kw) == (test_pos, extra_kw)
    # correct type
    assert pseudo.to_pseudo_tuple(test_pos, **extra_kw) == (test_pos, extra_kw)
    # kwargs
    assert pseudo.to_pseudo_tuple(pseudo1=1, pseudo2=2, pseudo3=3, **extra_kw) == \
        (test_pos, extra_kw)

    # too many positional arguments
    with pytest.raises(ValueError):
        pseudo.to_pseudo_tuple(1, 2, 3, 4)
    # valid kwargs, but passing in args too
    with pytest.raises(ValueError):
        pseudo.to_pseudo_tuple(1, pseudo1=1, pseudo2=2, pseudo3=3)


def test_multi_sequential():
    pseudo = Pseudo3x3('', name='mypseudo', concurrent=False)
    pseudo.wait_for_connection()

    assert pseudo.egu == 'a, b, c'

    pos2 = pseudo.PseudoPosition(pseudo1=0, pseudo2=0, pseudo3=0)
    pseudo.set(pos2, wait=True)
    time.sleep(1.0)
    pos1 = pseudo.PseudoPosition(pseudo1=.1, pseudo2=.2, pseudo3=.3)
    pseudo.set(pos1, wait=True)

    pseudo.real1.set(0, wait=True)
    pseudo.real2.set(0, wait=True)
    pseudo.real3.set(0, wait=True)

    pseudo.pseudo1.stop()

    pseudo.real3.set(0, wait=True)


def test_faulty_stopper():
    pseudo = FaultyPseudo1x3('', name='mypseudo', concurrent=False)
    pseudo.wait_for_connection()

    with pytest.raises(ExceptionBundle):
        # smoke-testing for coverage
        pseudo.pseudo1.stop()


def test_limits():
    pseudo = Pseudo3x3('', name='mypseudo', concurrent=True)
    assert pseudo.limits == ((-10, 10), (-10, 10), (0, 0))
    assert pseudo.low_limit == (-10, -10, 0)
    assert pseudo.high_limit == (10, 10, 0)


def test_read_describe():
    pseudo = Pseudo3x3('', name='mypseudo', concurrent=True)
    pseudo.wait_for_connection()
    desc_dict = pseudo.describe()
    desc_keys = ['source', 'upper_ctrl_limit', 'lower_ctrl_limit', 'shape',
                 'dtype', 'units']

    for key in desc_keys:
        assert key in desc_dict['mypseudo_pseudo3']

    read_dict = pseudo.read()
    read_keys = ['value', 'timestamp']
    for key in read_keys:
        assert key in read_dict['mypseudo_pseudo3']

    assert pseudo.read().keys() == pseudo.describe().keys()


def test_multi_concurrent():
    def done(**kwargs):
        logger.debug('** Finished moving (%s)', kwargs)
    pseudo = Pseudo3x3('', name='mypseudo', concurrent=True,
                       settle_time=0.1, timeout=25.0)
    assert pseudo.sequential is False
    assert pseudo.concurrent is True
    assert pseudo.settle_time == 0.1
    assert pseudo.timeout == 25.0
    pseudo.wait_for_connection()
    assert pseudo.connected
    assert tuple(pseudo.pseudo_positioners) == (pseudo.pseudo1, pseudo.pseudo2, pseudo.pseudo3)
    assert tuple(pseudo.real_positioners) == (pseudo.real1, pseudo.real2, pseudo.real3)

    logger.info('Move to (.2, .2, .2), which is (-.2, -.2, -.2) for real '
                'motors')
    pseudo.set(pseudo.PseudoPosition(.2, .2, .2), wait=True)
    logger.info('Position is: %s (moving=%s)', pseudo.position,
                pseudo.moving)
    pseudo.check_value((2, 2, 2))
    pseudo.check_value(pseudo.PseudoPosition(2, 2, 2))
    try:
        pseudo.check_value((2, 2, 2, 3))
    except ValueError as ex:
        logger.info('Check value failed, as expected (%s)', ex)
    real1 = pseudo.real1
    pseudo1 = pseudo.pseudo1
    try:
        pseudo.check_value((real1.high_limit + 1, 2, 2))
    except ValueError as ex:
        logger.info('Check value failed, as expected (%s)', ex)

    ret = pseudo.set((2, 2, 2), wait=False, moved_cb=done)
    assert ret.settle_time == 0.1
    count = 0
    while not ret.done:
        logger.info('Pos=%s %s (err=%s)', pseudo.position, ret, ret.error)
        count += 1
        if count > 1000:
            raise Exception
        time.sleep(0.1)
    logger.info('Single pseudo axis: %s', pseudo1)

    pseudo1.set(0, wait=True, timeout=5)
    assert pseudo1.target == 0
    pseudo1.sync()
    assert pseudo1.target == pseudo1.position
    # coverage
    pseudo1._started_moving

    try:
        pseudo1.check_value(real1.high_limit + 1)
    except ValueError as ex:
        logger.info('Check value for single failed, as expected (%s)', ex)

    logger.info('Move pseudo1 to 0, position=%s', pseudo.position)
    logger.info('pseudo1 = %s', pseudo1.position)

    def single_sub(**kwargs):
        # logger.info('Single sub: %s', kwargs)
        pass

    pseudo1.subscribe(single_sub, pseudo1.SUB_READBACK)
    ret = pseudo1.set(1, wait=False)
    assert pseudo.timeout == ret.timeout
    count = 0
    while not ret.done:
        logger.info('pseudo1.pos=%s Pos=%s %s (err=%s)', pseudo1.position,
                    pseudo.position, ret, ret.error)
        count += 1
        if count > 20:
            raise Exception
        time.sleep(0.1)

    logger.info('pseudo1.pos=%s Pos=%s %s (err=%s)', pseudo1.position,
                pseudo.position, ret, ret.error)
    copy(pseudo)
    pseudo.read()
    pseudo.describe()
    pseudo.read_configuration()
    pseudo.describe_configuration()
    repr(pseudo)
    str(pseudo)
    pseudo.pseudo1.read()
    pseudo.pseudo1.describe()
    pseudo.pseudo1.read_configuration()
    pseudo.pseudo1.describe_configuration()


def test_single_pseudo():
    logger.info('------- Sequential, single pseudo positioner')
    pos = Pseudo1x3('', name='mypseudo', concurrent=False)
    pos.wait_for_connection()
    reals = pos._real

    logger.info('Move to .2, which is (-.2, -.2, -.2) for real motors')
    pos.set((.2, ), wait=True)
    logger.info('Position is: %s (moving=%s)', pos.position, pos.moving)
    logger.info('Real positions: %s', [real.position for real in reals])

    logger.info('Move to -.2, which is (.2, .2, .2) for real motors')
    pos.set((-.2, ), wait=True)
    logger.info('Position is: %s (moving=%s)', pos.position, pos.moving)
    logger.info('Real positions: %s', [real.position for real in reals])

    copy(pos)
    pos.read()
    pos.describe()
    repr(pos)
    str(pos)


@pytest.mark.parametrize(
    'inpargs,inpkwargs,expected_position,expected_kwargs',
    [((1, 2, 3), {}, (1, 2, 3), {}),
     ((1, 2, ), {}, (1, 2, -3), {}),
     ((1, ), {}, (1, -2, -3), {}),

     (((1, 2, 3),), {}, (1, 2, 3), {}),
     (([1, 2, ],), {}, (1, 2, -3), {}),
     (((1, ),), {}, (1, -2, -3), {}),

     ((), {'pseudo1': 1, 'pseudo2': 2, 'pseudo3': 3}, (1, 2, 3), {}),
     ((), {'pseudo1': 1, 'pseudo2': 2}, (1, 2, -3), {}),
     ((), {'pseudo1': 1}, (1, -2, -3), {}),
     ((), {'pseudo1': 1, 'wait': True}, (1, -2, -3), {'wait': True}),

     (({'pseudo1': 1, 'pseudo2': 2, 'pseudo3': 3},), {}, (1, 2, 3), {}),
     (({'pseudo1': 1, 'pseudo2': 2},), {}, (1, 2, -3), {}),
     (({'pseudo1': 1},), {}, (1, -2, -3), {}),
     (({'pseudo1': 1, 'wait': True},), {'timeout': None},
      (1, -2, -3), {'wait': True, 'timeout': None}),


     ((1, 2, 3), {'timeout': 1}, (1, 2, 3), {'timeout': 1}),
     (((1, 2, 3),), {'timeout': 1}, (1, 2, 3), {'timeout': 1}),
     ]
)
def test_pseudo_position_input_3x3(hw, inpargs, inpkwargs,
                                   expected_position, expected_kwargs):
    pseudo3x3 = hw.pseudo3x3
    pseudo3x3.real1.set(1)
    pseudo3x3.real2.set(2)
    pseudo3x3.real3.set(3)

    out, extra_kwargs = pseudo3x3.to_pseudo_tuple(*inpargs, **inpkwargs)
    assert out == pseudo3x3.PseudoPosition(*expected_position)
    assert extra_kwargs == expected_kwargs

    pseudo3x3.set(*inpargs, **inpkwargs)
    assert pseudo3x3.position == pseudo3x3.PseudoPosition(*expected_position)


@pytest.mark.parametrize(
    'inpargs,inpkwargs',
    [((1, 2, 3, 5), {}),
     ((1, 2, 3), {'pseudo1': 1}),
     ((1, 2, 3), {'pseudo2': 1}),
     ((1, ), {'pseudo2': 1,
              'pseudo3': 1}),
     ((1, 2, ), {'pseudo3': 1}),
     ]
)
def test_pseudo_position_fail_3x3(hw, inpargs, inpkwargs):
    pseudo3x3 = hw.pseudo3x3
    with pytest.raises(ValueError):
        pseudo3x3.to_pseudo_tuple(*inpargs, **inpkwargs)


@pytest.mark.parametrize(
    'inpargs,inpkwargs,expected_position,expected_kwargs',
    [((1, 2, 3), {}, (1, 2, 3), {}),
     ((1, 2, ), {}, (1, 2, 3), {}),
     ((1, ), {}, (1, 2, 3), {}),

     (((1, 2, 3),), {}, (1, 2, 3), {}),
     (([1, 2, ],), {}, (1, 2, 3), {}),
     (((1, ),), {}, (1, 2, 3), {}),

     ((), {'real1': 1, 'real2': 2, 'real3': 3}, (1, 2, 3), {}),
     ((), {'real1': 1, 'real2': 2}, (1, 2, 3), {}),
     ((), {'real1': 1}, (1, 2, 3), {}),
     ((), {'real1': 1, 'foo': 'bar'}, (1, 2, 3), {'foo': 'bar'}),

     (({'real1': 1, 'real2': 2, 'real3': 3},), {}, (1, 2, 3), {}),
     (({'real1': 1, 'real2': 2},), {}, (1, 2, 3), {}),
     (({'real1': 1},), {}, (1, 2, 3), {}),
     (({'real1': 1, 'foo': 'bar'},), {'baz': 'buz'},
      (1, 2, 3), {'foo': 'bar', 'baz': 'buz'}),

     ((1, 2, 3), {'foo': 'bar'}, (1, 2, 3), {'foo': 'bar'}),
     ]
)
def test_real_position_input_3x3(hw, inpargs, inpkwargs,
                                 expected_position, expected_kwargs):
    pseudo3x3 = hw.pseudo3x3
    pseudo3x3.real1.set(1)
    pseudo3x3.real2.set(2)
    pseudo3x3.real3.set(3)

    out, extra_kwargs = pseudo3x3.to_real_tuple(*inpargs, **inpkwargs)
    assert out == pseudo3x3.RealPosition(*expected_position)
    assert extra_kwargs == expected_kwargs


@pytest.mark.parametrize(
    'inpargs,inpkwargs',
    [((1, 2, 3, 5), {}),
     ((1, 2, 3), {'real1': 1}),
     ((1, 2, 3), {'real2': 1}),
     ((1, ), {'real2': 1,
              'real3': 1}),
     ((1, 2, ), {'real3': 1}),
     (({'real3': 1, 'foo': 'bar'},), {'foo': 'bizz'}),
     ((), {})
     ]
)
def test_real_position_fail_3x3(hw, inpargs, inpkwargs):
    pseudo3x3 = hw.pseudo3x3
    with pytest.raises(ValueError):
        pseudo3x3.to_real_tuple(*inpargs, **inpkwargs)


def test_single_pseudo_with_sim(hw):
    logger.info('------- Sequential, single pseudo positioner')
    pos = hw.pseudo1x3

    reals = pos._real

    logger.info('Move to .2, which is (-.2, -.2, -.2) for real motors')
    pos.set((.2, ), wait=True)
    logger.info('Position is: %s (moving=%s)', pos.position, pos.moving)
    logger.info('Real positions: %s', [real.position for real in reals])

    logger.info('Move to -.2, which is (.2, .2, .2) for real motors')
    pos.set((-.2, ), wait=True)
    logger.info('Position is: %s (moving=%s)', pos.position, pos.moving)
    logger.info('Real positions: %s', [real.position for real in reals])

    copy(pos)
    pos.read()
    pos.describe()
    repr(pos)
    str(pos)


@pytest.mark.parametrize('typ', ('to_real_tuple', 'to_pseudo_tuple'))
@pytest.mark.parametrize('op', (operator.sub, operator.add))
@pytest.mark.parametrize('a,b', [
    ((0, 0, 0), (1, 1, 1)),
    ((1, 0, 1), (1, 1, 1)),
    ((9, 0, .3), (.3, .1, .5))
])
def test_pseudo_math(hw, a, b, op, typ):
    pos = hw.pseudo3x3
    a, _ = getattr(pos, typ)(a)
    b, _ = getattr(pos, typ)(b)

    # TODO switch to np asserts
    expected = op(np.asarray(a), np.asarray(b))
    assert (np.asarray(op(a, b)) == expected).all()
    assert (np.asarray(op(a, tuple(b))) == expected).all()
    assert (np.asarray(op(a, list(b))) == expected).all()
    assert (np.asarray(op(a, b._asdict())) == expected).all()
    assert (np.asarray(op(a, {})) == a).all()
    assert abs(op(a, b)) == np.sqrt(np.sum(expected ** 2))


def test_pseudo_hints(hw):
    pos = hw.pseudo3x3

    for j in (1, 2, 3):
        p = getattr(pos, 'pseudo{}'.format(j))
        assert p.hints['fields'] == [p.readback.name]
        p.readback.name = 'aardvark{}'.format(j)
        assert p.hints['fields'] == [p.readback.name]

    expected_fields = [getattr(pos, 'pseudo{}'.format(j)).readback.name
                       for j in (1, 2, 3)]
    assert pos.hints['fields'] == expected_fields
