# !/usr/bin/env python
#  -*- coding: utf-8 -*-
import numpy as np
import random


class StimuliPool(object):
    ''' where sequence of stimuli are generated by certain distribution,
        the distribution can be reimplemented in sub-class'''
    def __init__(self, N=None, S=None):
        if not N and not S:
            return None
        self.S = S  # the number of stimulus (or target propagation)
        self.N = N  # the number neurons
        whole_range = range(0, N)
        whole_set = set(whole_range)
        x_number_range = range(1, (N + 1) // 2)
        self.data = []
        for _ in range(S):
            n_x = random.choice(x_number_range)
            x_set = set(random.sample(whole_range, n_x))
            # to prevent large propgation
            n_y = np.min((n_x, random.choice(range(1, N - n_x))))
            to_propagate_set = set(random.sample(whole_set - x_set, n_y))
            y_set = x_set.union(to_propagate_set)
            self.data.append([x_set, [y_set]])
            # print x_set, y_set

    @classmethod
    def case_x_y_seperated(cls):
        stimu_pool = StimuliPool()
        stimu_pool.data = [
            [set([0, 1]), [set([2, 3])]],
            [set([4, 5]), [set([6, 7])]]]
        return stimu_pool

    @classmethod
    def case_x_overlapped(cls):
        stimu_pool = StimuliPool()
        stimu_pool.data = [
            [set([0, 1]), [set([2, 3])]],
            [set([1, 3]), [set([4, 5])]]]
        return stimu_pool

    @classmethod
    def case_y_overlapped(cls):
        stimu_pool = StimuliPool()
        stimu_pool.data = [
            [set([0, 1]), [set([2, 3])]],
            [set([4, 5]), [set([3, 6])]]]
        return stimu_pool

    @classmethod
    def case_x_y_overlapped(cls):
        stimu_pool = StimuliPool()
        stimu_pool.data = [
            [set([0, 1]), [set([2, 5])]],
            [set([1, 3]), [set([5, 6])]]]
        return stimu_pool

    def info(self):
        for i in self.data:
            print(i)


if __name__ == "__main__":
    pool = StimuliPool(20, 10)
    pool.info()
    print(StimuliPool.case_x_y_seperated().data)
    print(StimuliPool.case_x_overlapped().data)
    print(StimuliPool.case_y_overlapped().data)
    print(StimuliPool.case_x_y_overlapped().data)

