from __future__ import print_function, division

import os
import sys

import numpy as np
import pandas as pd
import sklearn.linear_model
import math
import pickle

import aiwolfpy
import aiwolfpy.cash

predictor = aiwolfpy.cash.Predictor_15()
def game_data_filter(df, day, phase='daily_initialize', agent=0):


    y = np.zeros(5460)
    # werewolves, possessed
    werewolves = []
    for i in range(1, 16):
        role = df["text"][i - 1].split()[2]
        if role == "WEREWOLF":
            werewolves.append(i)
        elif role == "POSSESSED":
            possessed = i

    for i in range(5460):
        if predictor.case15.case5460_df["agent_"+str(possessed)][i] == 2:
            if predictor.case15.case5460_df["agent_"+str(werewolves[0])][i] == 1:
                if predictor.case15.case5460_df["agent_"+str(werewolves[1])][i] == 1:
                    if predictor.case15.case5460_df["agent_"+str(werewolves[2])][i] == 1:
                        y[i] = 1

    # role
    role = "VILLAGER"
    if agent > 0:
        role = df["text"][agent - 1].split()[2]

    # filter by role
    if role in ["VILLAGER", "POSSESSED"]:
        df = df[df["type"].isin(["talk", "vote", "execute", "dead"])]
    elif role == "MEDIUM":
        df = df[df["type"].isin(["talk", "vote", "execute", "dead", "identify"])]
    elif role == "SEER":
        df = df[df["type"].isin(["talk", "vote", "execute", "dead", "divine"])]
    elif role == "BODYGUARD":
        df = df[df["type"].isin(["talk", "vote", "execute", "dead", "guard"])]
    elif role == "WEREWOLF":
        df = df[df["type"].isin(["talk", "vote", "execute", "dead", "whisper", "attack", "attack_vote"])]


    # agent
    if agent == 0:
        agent = 1

        # filter by time
    if phase == 'daily_initialize':
        df = df[df["day"] < day]
    else:
        df = df[(df["day"] < day) | ((df["day"] == day) & (df["type"] == 'talk'))]

    predictor.initialize({"agent":agent, "roleMap":{str(agent):role},"agentIdx":agent}, {})
    predictor.update_features(df.reset_index())
    predictor.update_df()

    return predictor.df_pred, y

count = 0

model = sklearn.linear_model.LogisticRegression()
wf = open("/home/AIWolfPy/train/train-log2.txt",mode="a")
for k in range(21,1000):
    x_1000 = np.zeros((5460000,60))
    y_1000 = np.zeros(5460000)
    ind = 0
    for i in range(100):
        log_path = "/home/AIWolfPy/pre_logs/gat2017_15/"+ "{0:03d}".format(k) + "/" + "{0:03d}".format(i) + ".log"
        for d in range(5):
            try:
                text_ = aiwolfpy.read_log(log_path)
            except:
                continue
            x, y = game_data_filter(text_, day=d, phase='vote')
            x_1000[(ind*5460):((ind+1)*5460), :] = x
            y_1000[(ind*5460):((ind+1)*5460)] = y
            ind += 1
            x, y = game_data_filter(text_, day=d+1, phase='daily_initialize')
            x_1000[(ind*5460):((ind+1)*5460), :] = x
            y_1000[(ind*5460):((ind+1)*5460)] = y
            ind += 1
        model.fit(x_1000, y_1000)
        text = "epoch :" + str(count) + " score :" + str(model.score(x_1000,y_1000))
        print(text)
        wf = open("/home/AIWolfPy/train/train-log.txt",mode="a")
        wf.write(text)
        wf.write("\n")
        wf.close()
        count+=1
    if k%10==0:
        filename = "/home/AIWolfPy/train/trainning_model-" + str(k) + ".sav"
        pickle.dump(model,open(filename,mode='wb'))
