from dzdy.abmodel import Agent, NetworkSet, TraitSet
import pandas as pd
from collections import OrderedDict


class Breeder:
    def __init__(self, sts, prefix='Ag'):
        self.Last = 0
        self.Prefix = prefix
        self.States = sts
        self.Traits = TraitSet()

    def breed(self, st, n=1, info=None):
        try:
            st = self.States[st]
        except KeyError:
            raise KeyError('The state is not well-defined')
        info = info if info else dict()
        start = self.Last + 1
        ags = [Agent('{}{}'.format(self.Prefix, start + i), st) for i in range(n)]
        for ag in ags:
            ag.update_info(info)
            self.Traits(ag.Info)

        self.Last += n
        return ags

    def append_trait(self, trait):
        self.Traits.append(trait)


class Population:
    def __init__(self, core, prefix='Ag'):
        self.Eve = Breeder(core.get_state_space(), prefix)
        self.Agents = OrderedDict()
        self.Networks = NetworkSet()

    def __getitem__(self, item):
        try:
            return self.Agents[item]
        except KeyError:
            raise KeyError('No this agent')

    def append_trait(self, fi):
        self.Eve.Traits.append(fi)

    def add_network(self, net):
        self.Networks[net.Name] = net

    def count(self, st=None):
        if st:
            if isinstance(st, str):
                st = self.Eve.States[st]
            return sum(st in ag for ag in self.Agents.values())
        else:
            return len(self.Agents)

    def neighbours(self, ag, net='|'):
        if net == '|':
            return self.Networks.neighbours_of(ag)
        elif net == '*':
            return self.Networks.neighbour_set_of(ag)
        else:
            try:
                return self.Networks.neighbours_of(ag, net)
            except KeyError:
                raise KeyError('No this net')

    def count_neighbours(self, ag, st=None, net=None):
        nes = self.neighbours(ag, net=net)
        if st:
            return sum(st in nei for nei in nes)
        else:
            return len(list(nes))

    def get_info_table(self):
        """
        get a pd.DataFrame table of the information of agents
        :return: pd.DataFrame table
        """
        return pd.DataFrame.from_records([ag.Info for ag in self.Agents.values()])

    def add_agent(self, atr, n=1, info=None):
        """

        :param atr: either the name of the attribute or entity of the attribute
        :param n: the number of agents need to be generate
        :param info: a dict of info. eg. {'Sex': 'F', 'Age': 26 }
        :return: a list of generated agents
        """
        ags = self.Eve.breed(atr, n, info)
        ags = list(ags)
        for ag in ags:
            self.Agents[ag.Name] = ag
            self.Networks.add_agent(ag)
        return ags

    def remove_agent(self, i):
        """
        remove agent from population, disconnecting all social links around it
        :param i: string, name of agent
        :return: deleted agent
        """
        try:
            ag = self[i]
            del self.Agents[i]
            self.Networks.remove_agent(ag)
        except KeyError:
            raise KeyError('No this agent')
        return ag

    def reform(self, net=None):
        self.Networks.reform(net)

    def first(self, n=5):
        for k, v in self.Agents.items():
            if n > 0:
                yield v
            else:
                return
            n -= 1

    def __str__(self):
        return "Population Size: {}\nNetwork: \n{}".format(len(self.Agents), repr(self.Networks))

    def __repr__(self):
        return "Population Size: {}, Network: {}".format(len(self.Agents), str(self.Networks))

@@ -1309,32 +1309,184 @@
 self, st=None):%0A
+        %22%22%22%0A        count agents with a certain state%0A        Args:%0A            st: target state or name of state%0A%0A        Returns: number%0A%0A        %22%22%22%0A
         if st:%0A 
@@ -2358,23 +2358,23 @@
         
-:r
+R
 eturn
+s
 : pd.Dat
@@ -2386,16 +2386,17 @@
 e table%0A
+%0A
         
@@ -2541,25 +2541,24 @@
     %22%22%22%0A
-%0A
         
 :param a
@@ -2549,22 +2549,52 @@
         
-:param
+Add agents%0A        Args:%0A           
  atr: ei
@@ -2659,22 +2659,19 @@
         
-:param
+   
  n: the 
@@ -2715,22 +2715,19 @@
         
-:param
+   
  info: a
@@ -2771,10 +2771,10 @@
 : 26
- 
 %7D
+%0A
 %0A   
@@ -2778,23 +2778,23 @@
         
-:r
+R
 eturn
+s
 : a list
@@ -2806,32 +2806,33 @@
 enerated agents%0A
+%0A
         %22%22%22%0A    
@@ -3066,17 +3066,17 @@
         
-r
+R
 emove ag
@@ -3149,14 +3149,25 @@
     
-:param
+Args:%0A           
  i: 
@@ -3188,24 +3188,25 @@
 f agent%0A
+%0A
         
 :return:
@@ -3201,16 +3201,20 @@
     
-:r
+R
 eturn
-:
+s: the
  del
@@ -3224,16 +3224,17 @@
 d agent%0A
+%0A
         

add docstring