akinator/akinator.py at main · StructCE/akinator · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def chooseLabel(df, labels):
  mins = [min(df[l].value_counts()) for l in labels]
  return labels[mins.index(max(mins))]

def genBinaryTree(df, labels):
  if len(df) == 1:
    return df['Name'].tolist()[0]
  if len(labels) == 1:
    return df["Name"].tolist()
  if len(df) == 0 or len(labels) == 0:
    return None

  label = chooseLabel(df, labels)
  newLabels = [l for l in labels if l!= label]

  return {
    "label": label,
    -1: genBinaryTree(df[df[label] != 1], newLabels),
    1: genBinaryTree(df[df[label] != -1], newLabels),
  }

class Akinator:
  df = None
  tree = {}
  executions = {}
  lastId = 0

  def __init__(self, dataFrame):
    self.df = dataFrame

  def genId(self):
    id = self.lastId
    self.lastId += 1
    return id

  def createExecution(self):
    id = self.genId()
    self.executions[id] = [self.tree]
    return id

  def genTreeFromDf(self):
    self.tree = self.genTree(self.df, self.df.columns.tolist()[2:])

  def genTree(self, df, labels):
    if len(df) == 1:
      return df['Name'].tolist()[0]
    if len(labels) == 1:
      return df["Name"].tolist()
    if len(df) == 0 or len(labels) == 0:
      return None

    label = chooseLabel(df, labels)
    newLabels = [l for l in labels if l!= label]

    return {
      "label": label,
      -1: genBinaryTree(df[df[label] != 1], newLabels),
      1: genBinaryTree(df[df[label] != -1], newLabels),
    }

  def addCharacter(character):
    # A implementar
    pass

  def getQuestion(self, id):
    if not id in self.executions:
      return None
    if len(self.executions[id]) == 0:
      return None
    return self.executions[id][-1]['label']

  def execute(self, id, inp):
    if self.executions[id] == []:
      return None
    tree = self.executions[id].pop()

    def evaluate(inp):
      if type(tree[inp]) == str or type(tree[inp]) == list:
        return tree[inp]
      else:
        self.executions[id].append(tree[inp])
        return None

    if inp != 0:
       return evaluate(inp)
    else:
      result = evaluate(-1)
      if result!= None:
        return result
      return evaluate(1)