from fireworks import message as messi
from fireworks import Message, TensorMessage
import torch
import numpy as np
from itertools import product
import pandas as pd
from itertools import count

tensors = {
    'a': torch.Tensor([1,2,3]),
    'b': torch.Tensor([4,5,6]),
    }
vectors = {
    'c': np.array([7,8,9]),
    'd': np.array([10,11,12]),
}
dtensors = {
    'a': torch.Tensor([[1,2,3],[4,5,6],[7,8,9]]),
    'b': torch.Tensor([[-1,-2,-3],[-4,-5,-6], [-7,-8,-9]]),
}

def test_compute_length():

    l = messi.compute_length(tensors)
    assert l == 3
    l = messi.compute_length(vectors)
    assert l == 3

def test_extract_tensors():

    target = {**tensors, **vectors}
    t, v = messi.extract_tensors(target)
    assert t == tensors
    assert v == vectors
    t, v = messi.extract_tensors(tensors)
    assert t == tensors
    assert v == {}
    t, v = messi.extract_tensors(vectors)
    assert t == {}
    assert v == vectors

def test_complement():

    n = 10
    index = 7
    complement = messi.complement(index, n)
    assert complement == [0,1,2,3,4,5,6,8,9]
    index = slice(2,5)
    complement = messi.complement(index, n)
    assert complement == [0,1,5,6,7,8,9]
    index = [2,4,6]
    complement = messi.complement(index, n)
    assert complement == [0,1,3,5,7,8,9]

def test_Message():
    """ Test init, getitem, and len methopl. """

    def attribute_test(message, length = 3):
        assert len(message) == length
        assert message[0].tensors() == {
            'a': torch.Tensor([1]),
            'b': torch.Tensor([4]),
        }
        assert message[0].dataframe().equals(pd.DataFrame({
            'c': np.array([7]),
            'd': np.array([10]),
        }))

        assert message[0] == Message({'a': torch.Tensor([1]),'b': torch.Tensor([4])}, pd.DataFrame({'c': np.array([7]),'d': np.array([10]),}))

        assert message[1:3].tensors() == {
            'a': torch.Tensor([2,3]),
            'b': torch.Tensor([5,6]),
        }
        assert message[1:3].dataframe().equals(pd.DataFrame({
            'c': np.array([8,9]),
            'd': np.array([11,12]),
        }))

        assert (message['a'] == torch.Tensor([1,2,3])).all()
        assert message[['a','c']] == Message({'a': torch.Tensor([1,2,3]), 'c': np.array([7,8,9])})

        assert message[1:3] == Message({'a': torch.Tensor([2,3]),'b': torch.Tensor([5,6])}, pd.DataFrame({'c': np.array([8,9]),'d': np.array([11,12])}))

        # Test length
        assert len(message) == length
        # Test __getitem__

    # Init empty message
    m = Message()
    assert len(m) == 0
    # Init message from tensor_dict / TensorMessage and dict of arrays / dataframe using positional arguments.
    tensor_message = TensorMessage(tensors)
    tensor_as_message = Message(tensors = tensors)
    df = pd.DataFrame(vectors)
    df_as_message = Message(df = vectors)

    # Try every combination
    tensor_options = [tensors, tensor_message, tensor_as_message]
    vector_options = [vectors, df, df_as_message]

    for t, v in product(tensor_options, vector_options):
        m = Message(t, v)
        attribute_test(m)
        m = Message(tensors = t, df = v)
        attribute_test(m)

    # Test one sided Messages
    for t in tensor_options:
        m = Message(t, None)
        assert len(m) == 3
        assert m == Message(tensors)
    for v in vector_options:
        m = Message(None, v)
        assert len(m) == 3
        assert m == Message(vectors)

    # Init message from a single dict
    everything = {**tensors, **vectors}
    m = Message(everything)
    attribute_test(m)

def test_Message_from_objects():

    v = vectors.copy()
    t = tensors.copy()
    v['c'] = np.array([1.,2.])
    v['r'] = 'howdy'
    t['a'] = torch.randn(5)
    t['q'] = torch.randn([4,3])
    combined = {**t, **v}

    m = Message.from_objects(t, v)
    assert (set(m.keys()) == set(['c','d','r','b','a','q']))
    for key in ['c','d','b','a','q']:
        assert (m[key][0] == combined[key]).all()
    assert m['r'][0] == combined['r']
    assert len(m) == 1

def test_getitem():

    m = Message(tensors, vectors)
    assert m[0] == Message({'a': torch.Tensor([1]), 'b': torch.Tensor([4])}, {'c': np.array([7]), 'd': np.array([10])})
    assert m[[0,2]] == Message({'a': torch.Tensor([1,3]), 'b': torch.Tensor([4,6])}, {'c': np.array([7,9]), 'd': np.array([10,12])})
    # Check that out of bounds index calls raise errors
    try:
        m[3]
        assert False
    except IndexError:
        assert True

    try:
        m[3:5]
        assert False
    except IndexError:
        assert True

def test_cache(): pass

def test_tensors():
    m = Message(tensors, vectors)
    t = m.tensors()
    assert t == TensorMessage(tensors)
    t = m.tensors(keys=['a'])
    assert t == TensorMessage({'a': tensors['a']})
    t = m.tensors(keys=['a','c'])
    assert t == TensorMessage({'a': tensors['a'], 'c': torch.Tensor(vectors['c'])})

def test_df():
    m = Message(tensors, vectors)
    df = m.dataframe()
    assert df.equals(pd.DataFrame(vectors))
    df = m.dataframe(keys=['c'])
    assert df.equals(pd.DataFrame({'c': vectors['c']}))
    df = m.dataframe(keys=['c','a'])
    assert (df == (pd.DataFrame({'c': vectors['c'], 'a': np.array(tensors['a'])}))).all().all()

def test_to_dataframe():

    mo = Message(tensors,vectors)
    # no = mo.to_dataframe()
    # assert no.tensor_message == {}
    # assert (no['a'] == mo['a']).all()
    # assert (no['b'] == mo['b']).all()
    # for letter in ['a','b','c','d']:
    #     assert letter in no.df
    lo = Message(dtensors, vectors)
    ok = lo.to_dataframe()
    for i in range(3):
        assert (ok['a'][i] == dtensors['a'][i].numpy()).all()
        assert (ok['b'][i] == dtensors['b'][i].numpy()).all()

def test_cpu_gpu():
    m = Message(tensors, vectors)
    m.cpu()
    assert set(m.tensors().keys()) == set(['a','b'])
    for key, tensor in m.tensors().items():
        assert tensor.device.type == 'cpu'
    if torch.cuda.is_available():
        m.cuda()
        for key, tensor in m.tensors().items():
            assert tensor.device.type == 'cuda'
        m.cpu()
        for key, tensor in m.tensors().items():
            assert tensor.device.type == 'cpu'

def test_append():

    t = tensors
    v = vectors
    m1 = Message(t, v)
    m2 = Message(t, v)
    m3 = Message(t)
    m4 = TensorMessage(t)
    m5 = Message(pd.DataFrame(v))
    m6 = pd.DataFrame(v)

    m0 = Message()
    assert(len(m0) == 0)
    m = m0.append(Message(t))
    assert m == Message(t)
    m = m0.append(Message(v))
    assert m == Message(v)
    m = m0.append(Message(t,v))
    assert m == Message(t,v)

    m = m1.append(m2)
    assert len(m) == 6
    assert m == Message({'a': torch.Tensor([1,2,3,1,2,3]), 'b': torch.Tensor([4,5,6,4,5,6])}, {'c': np.array([7,8,9,7,8,9]), 'd': np.array([10,11,12,10,11,12])})
    m = m3.append(t)
    assert len(m) == 6
    assert m == Message({'a': torch.Tensor([1,2,3,1,2,3]), 'b': torch.Tensor([4,5,6,4,5,6])})
    m = m3.append(m3)
    assert len(m) == 6
    assert m == Message({'a': torch.Tensor([1,2,3,1,2,3]), 'b': torch.Tensor([4,5,6,4,5,6])})
    m = m3.append(m4)
    assert len(m) == 6
    assert m == Message({'a': torch.Tensor([1,2,3,1,2,3]), 'b': torch.Tensor([4,5,6,4,5,6])})

    m = m4.append(t)
    assert len(m) == 6
    assert m == TensorMessage({'a': torch.Tensor([1,2,3,1,2,3]), 'b': torch.Tensor([4,5,6,4,5,6])})
    m = m4.append(m3)
    assert len(m) == 6
    assert m == TensorMessage({'a': torch.Tensor([1,2,3,1,2,3]), 'b': torch.Tensor([4,5,6,4,5,6])})
    m = m4.append(m4)
    assert len(m) == 6
    assert m == TensorMessage({'a': torch.Tensor([1,2,3,1,2,3]), 'b': torch.Tensor([4,5,6,4,5,6])})

    m = m5.append(v)
    assert len(m) == 6
    assert m == Message({'c': np.array([7,8,9,7,8,9]), 'd': np.array([10,11,12,10,11,12])})
    m = m5.append(m5)
    assert len(m) == 6
    assert m == Message({'c': np.array([7,8,9,7,8,9]), 'd': np.array([10,11,12,10,11,12])})
    m = m5.append(m6)
    assert len(m) == 6
    assert m == Message({'c': np.array([7,8,9,7,8,9]), 'd': np.array([10,11,12,10,11,12])})

    # Test type conversions on appending to TensorMessage
    m = m4.append({'a': np.array([42]), 'b': np.array([24])})
    assert len(m) == 4
    assert m == TensorMessage({'a': torch.Tensor([1,2,3,42]), 'b': torch.Tensor([4,5,6,24])})


def test_join():

    t = tensors
    v = vectors
    t2 = {'d': torch.Tensor([13,14,15])}
    v2 = {'e': np.array([16,17,18])}
    m1 = Message(t,v)
    m2 = Message(t)
    m2_t = TensorMessage(t)
    m3 = Message(v)
    m4 = Message(t2,v2)
    m5 = Message(t2)
    m5_t = TensorMessage(t2)
    m6 = Message(v2)
    m7 = Message(t,v2)
    m8 = Message(t2, v)

    # Test if a tensor message can be merged into a message and vice versa
    assert m2.merge(m3) == m1
    assert m3.merge(m2) == m1
    assert m3.merge(m2_t) == m1
    assert m3.merge(t) == m1

    # Test if the tensors in messages can be merged
    assert m2.merge(t2) == Message({**t, **t2})
    assert m2.merge(m5) == Message({**t, **t2})
    assert m2.merge(m5_t) == Message({**t, **t2})
    assert m2_t.merge(t2) == TensorMessage({**t, **t2})
    assert m2_t.merge(m5) == TensorMessage({**t, **t2})
    assert m2_t.merge(m5_t) == TensorMessage({**t, **t2})

    # Test if the dataframes in messages can be merged
    assert m3.merge(m6) == Message({**v, **v2})
    assert m6.merge(m3) == Message({**v, **v2})
    assert m3.merge(v2) == Message({**v, **v2})

def test_Message_set_get():

    # Test point updates
    email = Message(tensors, vectors)
    gmail = Message({'a':torch.Tensor([1,42,3]), 'b':torch.Tensor([4,43,6]), 'c': np.array([7,99,9]), 'd': np.array([10,100,12])})
    replacement = {'a': torch.Tensor([42]), 'b': torch.Tensor([43]), 'c': np.array([99]), 'd': np.array([100])}
    assert len(email) == 3
    assert email != gmail
    email[1] = replacement
    assert email == gmail
    # Test ranged updates
    email = Message(tensors, vectors)
    gmail = Message({'a':torch.Tensor([1,42,33]), 'b':torch.Tensor([4,43,66]), 'c': np.array([7,99,99]), 'd': np.array([10,100,122])})
    replacement = {'a': torch.Tensor([42,33]), 'b': torch.Tensor([43,66]), 'c': np.array([99,99]), 'd': np.array([100,122])}
    assert email != gmail
    email[1:3] = replacement
    assert email == gmail
    # Test updates using lists as indexes
    email = Message(tensors, vectors)
    assert email != gmail
    email[[1,2]] = replacement
    assert email == gmail
    # Test column updates
    email['a'] = torch.Tensor([9,9,9])
    assert torch.equal(email['a'], torch.Tensor([9,9,9]))
    email['c'] = np.array([9,9,9])
    assert email['c'].equals(pd.Series([9,9,9]))
    # Test column updates that switch from df to tensor and vice-versa
    email = Message(tensors, vectors)
    assert set(email.columns) == set(['a','b','c','d'])
    assert set(email.tensor_message.columns) == set(['a','b'])
    assert set(email.df.columns) == set(['c','d'])
    new_a = np.array([1,2,3]) # Switch from tensor to vector
    email['a'] = new_a
    assert set(email.columns) == set(['a','b','c','d'])
    assert set(email.tensor_message.columns) == set(['b'])
    assert set(email.df.columns) == set(['a','c','d'])
    assert (email['a'] == new_a).all()
    new_c = torch.Tensor([7,8,9])
    email['c'] = new_c
    assert set(email.columns) == set(['a','b','c','d'])
    assert set(email.tensor_message.columns) == set(['b','c'])
    assert set(email.df.columns) == set(['a','d'])
    assert (email['c'] == new_c).all()
    # Test column updates that end up clearing either self.df or self.tensor_message
    email = Message(tensors, vectors)
    df = email.dataframe(['a', 'b'])
    assert len(email) == 3
    assert len(email.tensor_message) == 3
    assert len(email.df) == 3
    email[['a','b']] = df
    assert len(email) == 3
    assert len(email.tensor_message) == 0
    assert len(email.df) == 3
    # TODO: Test the other way around

def test_Message_del():

    t = {
        'a': torch.Tensor([1,2,3]),
        'b': torch.Tensor([4,5,6]),
        }
    v = {
        'c': np.array([7,8,9]),
        'd': np.array([10,11,12]),
    }
    t2 = {
        'a': torch.Tensor([1,2]),
        'b': torch.Tensor([4,5]),
        }
    v2 = {
        'c': np.array([7,8]),
        'd': np.array([10,11]),
    }
    t3 = {
        'a': torch.Tensor([1]),
        'b': torch.Tensor([4]),
        }
    v3 = {
        'c': np.array([7]),
        'd': np.array([10]),
    }

    # Test deletions for messages with only tensors, only df, and both
    # Test point deletions
    m = Message(t,v)
    m1 = Message(t)
    m2 = Message(v)
    assert m != Message(t2,v2)
    assert m1 != Message(t2)
    assert m2 != Message(v2)
    assert len(m) == 3
    assert len(m1) == 3
    assert len(m2) == 3
    del m[2]
    del m1[2]
    del m2[2]
    assert len(m) == 2
    assert len(m1) == 2
    assert len(m2) == 2
    assert m == Message(t2,v2)
    assert m1 == Message(t2)
    assert m2 == Message(v2)
    # Test range deletions
    m = Message(t,v)
    m1 = Message(t)
    m2 = Message(v)
    assert m != Message(t3,v3)
    assert m1 != Message(t3)
    assert m2 != Message(v3)
    assert len(m) == 3
    assert len(m1) == 3
    assert len(m2) == 3
    del m[1:3]
    del m1[1:3]
    del m2[1:3]
    assert len(m) == 1
    assert len(m1) == 1
    assert len(m2) == 1
    assert m == Message(t3,v3)
    assert m1 == Message(t3)
    assert m2 == Message(v3)
    # Test list deletions
    m = Message(t,v)
    m1 = Message(t)
    m2 = Message(v)
    assert m != Message(t3,v3)
    assert m1 != Message(t3)
    assert m2 != Message(v3)
    assert len(m) == 3
    assert len(m1) == 3
    assert len(m2) == 3
    del m[[1,2]]
    del m1[[1,2]]
    del m2[[1,2]]
    assert len(m) == 1
    assert len(m1) == 1
    assert len(m2) == 1
    assert m == Message(t3,v3)
    assert m1 == Message(t3)
    assert m2 == Message(v3)

    # Test column deletions
    m = Message(t,v)
    assert set(m.columns) == set(['a','b','c','d'])
    del m['a']
    assert set(m.columns)== set(['b','c','d'])
    del m['c']
    assert set(m.columns) == set(['b','d'])

def test_Message_iter():

    m = Message(tensors, vectors)
    l = len(m)
    for x,i in zip(m, count()):
        assert type(x) is Message
        if i > l:
            assert False
    assert i == l - 1

    t = TensorMessage(tensors)
    l = len(t)
    for x,i in zip(t, count()):
        assert type(x) is TensorMessage
        if i > l:
            assert False
    assert i == l - 1

def test_map(): pass

def test_TensorMessage():

    a = [1,2,3]
    b = [4, 5, 6]

    # Test init
    empty = TensorMessage()
    assert len(empty) == 0
    assert empty.keys() == {}.keys()
    email = TensorMessage({'a': a, 'b':b})

    #   Test error cases
    # TODO: test error cases

    # Test length
    assert len(email) == 3
    # Test getitem
    x = email[2]
    assert set(x.keys()) == set(['a','b'])
    assert (x['a'] == torch.Tensor([3])).all()
    assert (x['b'] == torch.Tensor([6])).all()
    x = email[0:2]
    assert set(x.keys()) == set(['a','b'])
    assert (x['a'] == torch.Tensor([1,2])).all()
    assert (x['b'] == torch.Tensor([4,5])).all()

    # Test for length 1 init
    gmail = TensorMessage({'a':1, 'b': 80})
    assert len(gmail) == 1
    y = gmail[0]
    assert set(y.keys()) == set(['a','b'])
    assert (y['a'] == torch.Tensor([1])).all()
    assert (y['b'] == torch.Tensor([80])).all()

    # Test extend
    yahoomail = email.append(gmail)
    assert len(yahoomail) == 4
    z = yahoomail[0:4]
    assert set(z.keys()) == set(['a','b'])
    assert (z['a'] == torch.Tensor([1,2,3,1])).all()
    assert (z['b'] == torch.Tensor([4, 5, 6, 80])).all()

def test_TensorMessage_set_get_del():

    a = [1,2,3]
    b = [4, 5, 6]
    email = TensorMessage({'a': a, 'b':b})
    replacement = {'a': torch.Tensor([42]), 'b': torch.Tensor([43])}
    gmail = TensorMessage({'a':torch.Tensor([42,2,3]), 'b':  torch.Tensor([43,5,6])})
    yahoomail = TensorMessage({'a':torch.Tensor([2,3]), 'b': torch.Tensor([5,6])})
    assert email != gmail
    email[0] = replacement
    assert email == gmail
    assert len(email) == 3
    email['a'] = torch.Tensor([9,9,9])
    assert torch.equal(email['a'], torch.Tensor([9,9,9]))
    assert gmail != yahoomail
    del gmail[0]
    assert len(gmail) == 2
    assert gmail == yahoomail
    # Test column deletions
    email = TensorMessage({'a': a, 'b':b})
    assert set(email.columns) == set(['a','b'])
    del email['a']
    assert set(email.columns) == set(['b'])
    # Test that out of bounds requests raise errors
    try:
        email[3]
        assert False
    except IndexError:
        assert True
    try:
        email[3:5]
        assert False
    except IndexError:
        assert True

    # Test length adjustment if all columns are deleted
    zohomail = TensorMessage({'a':a,'b':b})
    assert len(zohomail) == 3
    del zohomail['a']
    assert len(zohomail) == 3
    del zohomail['b']
    assert len(zohomail) == 0

def test_TensorMessage_eq():

    a = [1,2,3]
    b = [4, 5, 6]

    # Test init
    email = TensorMessage({'a': a, 'b':b})
    gmail = TensorMessage(email)

def test_cat():

    m = Message(tensors, vectors)
    m0 = m[0]
    m1 = m[1]
    m2 = m[2]
    babaghanush = messi.cat([m0,m1,m2])
    assert babaghanush == m

def test_TensorMessage_permute():
    a = [1,2,3]
    b = [4, 5, 6]
    email = TensorMessage({'a': a, 'b':b})
    gmail = email.permute([2,1,0])
    assert gmail == TensorMessage({'a':[3,2,1], 'b':[6,5,4]})
    gmail = email.permute([0,0,0])
    assert gmail == TensorMessage({'a':[1,1,1], 'b':[4,4,4]})

def test_permute():

    tensors = {
        'a': torch.Tensor([1,2,3]),
        'b': torch.Tensor([4,5,6]),
        }
    vectors = {
        'c': np.array([7,8,9]),
        'd': np.array([10,11,12]),
    }
    email = Message(tensors, vectors)
    gmail = email.permute([2,1,0])
    assert gmail == Message({'a':[3,2,1], 'b':[6,5,4]}, {'c': np.array([9,8,7]), 'd': np.array([12,11,10])})
    gmail = email.permute([0,0,0])
    assert gmail == Message({'a':[1,1,1], 'b':[4,4,4]}, {'c': np.array([7,7,7]), 'd': np.array([10,10,10])})
    # Test with only tensors
    email = Message(tensors)
    gmail = email.permute([2,1,0])
    assert gmail == Message({'a':[3,2,1], 'b':[6,5,4]}, {})
    gmail = email.permute([0,0,0])
    assert gmail == Message({'a':[1,1,1], 'b':[4,4,4]}, {})
    # Test with only dataframes
    email = Message(vectors)
    gmail = email.permute([2,1,0])
    assert gmail == Message({'c': np.array([9,8,7]), 'd': np.array([12,11,10])})
    gmail = email.permute([0,0,0])
    assert gmail == Message({'c': np.array([7,7,7]), 'd': np.array([10,10,10])})

def test_to_csv():
    m = Message(tensors, vectors)
    pass #TODO: Implement

def test_to_pickle():
    m = Message(tensors, vectors)
    pass #TODO: Implement

def test_to_sql():
    m = Message(tensors, vectors)
    pass #TODO: Implement

def test_to_dict():
    m = Message(tensors, vectors)
    md = m.to_dict()
    assert type(md) is dict
    assert (md['c'] == md['c'])
    assert (md['d'] == md['d'])
    assert (md['a'] == np.array(md['a'])).all()
    assert (md['b'] == np.array(md['b'])).all()

def test_to_excel():
    m = Message(tensors, vectors)
    pass #TODO: Implement

def test_to_json():
    m = Message(tensors, vectors)
    pass #TODO: Implement

def test_to_string():
    m = Message(tensors, vectors)
    pass #TODO: Implement
