
from datetime import datetime
from collections import defaultdict
from unittest.mock import Mock
import pytest

import neo
from pyNN import recording, errors
from pyNN.recording import Variable

# def test_rename_existing():

# def test_gather():
#import time
# for x in range(7):
#    N = pow(10, x)
#    local_data = np.empty((N,2))
#    local_data[:,0] = np.ones(N, dtype=float)*comm.rank
#    local_data[:,1] = np.random.rand(N)
#
#    start_time = time.time()
#    all_data = gather(local_data)
#    #print(comm.rank, "local", local_data)
#    if comm.rank == 0:
#    #    print("all", all_data)
#        print(N, time.time()-start_time)

# def test_gather_no_MPI():

# def test_gather_dict():

# def test_mpi_sum():


class MockState(object):

    def __init__(self, mpi_rank):
        self.mpi_rank = mpi_rank
        self.num_processes = 9
        self.dt = 0.123
        self.running = True
        self.recorders = set([])
        self.t = 0.0


class MockSimulator(object):
    name = "MockSimulator"

    def __init__(self, mpi_rank):
        self.state = MockState(mpi_rank)


class MockNeoIO(object):
    filename = "fake_file"
    write = Mock()


class MockRecorder(recording.Recorder):
    _simulator = MockSimulator(mpi_rank=0)

    def _get_current_segment(self, filter_ids=None, variables='all', clear=False):
        segment = Mock()
        segment.analogsignals = [Mock(), Mock()]
        segment.proxy_for = neo.Segment
        return segment

    def _localize_variables(self, variables, locations):
        # If variables is a single string, encapsulate it in a list.
        if isinstance(variables, str) and variables != 'all':
            variables = [variables]
        if isinstance(locations, str):
            locations = [locations]
        resolved_variables = []

        if locations is None:
            for var_path in variables:
                resolved_variables.append(recording.Variable(location=None, name=var_path, label=None))
        else:
            raise NotImplementedError

        return resolved_variables



class MockPopulation(object):
    size = 11
    first_id = 2454
    last_id = first_id + size
    label = "mock population"
    celltype = Mock(always_local=False)
    annotations = {'knights_say': 'Ni!'}

    def __len__(self):
        return self.size

    def can_record(self, variable, location=None):
        if variable in ["spikes", "v", "gsyn_exc", "gsyn_inh", "spam"]:
            return True
        else:
            return False

    def id_to_index(self, id):
        return id

    def describe(self):
        return "mock population"


class MockNeoBlock(object):

    def __init__(self):
        self.name = None
        self.description = None
        self.segments = [Mock()]
        self.rec_datetime = datetime.now()

    def annotate(self, **annotations):
        pass


def test_Recorder_create():
    p = MockPopulation()
    r = MockRecorder(p)
    assert r.population == p
    assert r.file == None
    assert r.recorded == defaultdict(set)


def test_Recorder_invalid_variable():
    p = MockPopulation()
    r = MockRecorder(p)
    all_ids = (MockID(0, True), MockID(1, False), MockID(
        2, True), MockID(3, True), MockID(4, False))
    with pytest.raises(errors.RecordingError):
        r.record('foo', all_ids)


class MockID(object):

    def __init__(self, id, local):
        self.id = id
        self.local = local


def test_record():
    p = MockPopulation()
    r = MockRecorder(p)
    r._record = Mock()
    spam_var = Variable(location=None, name='spam', label=None)
    assert r.recorded == defaultdict(set)

    all_ids = (MockID(0, True), MockID(1, False), MockID(
        2, True), MockID(3, True), MockID(4, False))
    first_ids = all_ids[0:3]
    r.record('spam', first_ids)
    assert r.recorded[spam_var] == set(id for id in first_ids if id.local)
    assert len(r.recorded[spam_var]) == 2
    r._record.assert_called_with(spam_var, r.recorded[spam_var], None)

    more_ids = all_ids[2:5]
    r.record('spam', more_ids)
    assert r.recorded[spam_var] == set(id for id in all_ids if id.local)
    assert len(r.recorded[spam_var]) == 3
    r._record.assert_called_with(spam_var, set(all_ids[3:4]), None)


def test_filter_recorded():
    p = MockPopulation()
    r = MockRecorder(p)
    r._record = Mock()
    spam_var = Variable(location=None, name='spam', label=None)
    spikes_var = Variable(location=None, name='spikes', label=None)
    all_ids = (MockID(0, True), MockID(1, False), MockID(2, True), MockID(3, True), MockID(4, False))
    r.record(['spikes', 'spam'], all_ids)
    assert r.recorded[spikes_var] == set(id for id in all_ids if id.local)
    assert r.recorded[spam_var] == set(id for id in all_ids if id.local)

    filter = all_ids[::2]
    filtered_ids = r.filter_recorded(spam_var, filter)
    assert filtered_ids == set(id for id in filter if id.local)

    assert r.filter_recorded(spikes_var, None) == r.recorded[spikes_var]


def test_get():
    p = MockPopulation()
    r = MockRecorder(p)
    data = r.get('spikes')
    assert data.name == p.label
    assert data.description == p.describe()

# def test_write__with_filename__compatible_output__gather__onroot():
#    orig_metadata = recording.Recorder.metadata
#    #recording.Recorder.metadata = {'a': 2, 'b':3}
#    p = MockPopulation()
#    r = MockRecorder(p)
#    #fake_data = MockNeoBlock()
#    r._get_current_segment = Mock() #return_value=fake_data)
#    output_io = MockNeoIO()
#    r.write("spikes", file=output_io, gather=True)
#    #recording.Recorder.metadata = orig_metadata
#    output_io.write.assert_called_with(fake_data)


def test_metadata_property():
    p = MockPopulation()
    r = MockRecorder(p)
    assert r.metadata == {
        'first_id': 2454, 'label': 'mock population',
        'dt': 0.123, 'last_id': 2465, 'size': 11,
        'first_index': 0, 'last_index': 11, 'knights_say': 'Ni!',
        'simulator': 'MockSimulator', 'mpi_processes': 9}


# def test_count__spikes_gather():

# def test_count__spikes_nogather():


# def test_count__other():
