from env.security_env import SecurityEnv, import_security_env
import numpy as np


class ManualPolicy(object):
    def __init__(self, senv: SecurityEnv, side):
        self.cache = dict()
        self.ob_encoder = senv.get_ob_encoders()[side]
        self.ob_namer = senv.get_ob_namers()[side]

    def get_strategy(self):
        def strategy(ob):
            eob = self.ob_encoder(ob)
            if eob in self.cache:
                return self.cache[eob]
            else:
                s = input(self.ob_namer(ob) + ": ")
                res = list(map(float, s.split()))
                # print(res)
                self.cache[eob] = np.array(res) / np.sum(np.array(res))
                return self.cache[eob]
        return strategy


if __name__ == "__main__":
    # n_slots = 2
    # n_types = 2
    # prior = [.3, .7]
    # n_rounds = 2
    # env = SecurityEnv(n_slots=n_slots, n_types=n_types, prior=prior, n_rounds=n_rounds)
    env = import_security_env("../result/setting.pkl")
    print(env.payoff)

    # def strategy(ob):
    #     return np.ones(shape=n_slots) / n_slots

    # policy = ManualPolicy(env, 0)
    # print(env.calc_exploitability(0, policy.get_strategy()))

    policy = ManualPolicy(env, 1)
    print(env.calc_exploitability(1, policy.get_strategy()))
