import numpy as np

from modular import Data
from helpers import paths, dist_funcs
from opt_cost import OptimalWeights
from KLFunctions import dBernoulli


class Problem:
    def __init__(self, rewardsdist, costsdist, costs, rewards, delta, N):
        self.rewardsdist = rewardsdist
        self.costsdist = costsdist
        self.costs = costs
        self.rewards = rewards
        self.delta = delta
        self.N = N

    def get_data(self, algo):
        costs = paths(self.costsdist, self.costs, self.N)
        rewards = paths(self.rewardsdist, self.rewards, self.N)
        data = Data(algo, costs, rewards, self.rewardsdist, delta=self.delta)
        return data

    def get_name(self):
        filename = "{}_rewards-{}_costs-mu_{}-costs_{}-{}_samples".format(self.rewardsdist,
                                                                          self.costsdist,
                                                                          self.rewards,
                                                                          self.costs, self.N)
        return filename

    def lower_bound(self):
        d, _, __ = dist_funcs(self.rewardsdist)
        T, w_vec = OptimalWeights(d, self.rewards, self.costs, delta=self.delta)
        lower_bound = dBernoulli(self.delta, 1 - self.delta)*T
        return lower_bound


N = 100

tas_mu1 = [0.5, 0.45, 0.43, 0.4]
tas_mu2 = [0.3, 0.21, 0.2, 0.19, 0.18]
tas_delta = 0.1
delta = 1e-9

tas_1 = Problem("Gaussian", "Deterministic",
                [1, 1, 1, 1], tas_mu1, tas_delta, N)

tas_2 = Problem("Gaussian", "Deterministic",
                [10, .8, .6, .4, .2], tas_mu2, tas_delta, N)

gaussian_deterministic = Problem("Gaussian", "Deterministic",
                           [1, .1, .01], [.9, .6, .3], delta, N)

gaussian_default = Problem("Gaussian", "Bernoulli",
                           [1, .1, .01], [1.5, 1.0, 0.5], delta, N)

gaussian_4 = Problem("Gaussian", "Bernoulli",
                           [10, 1, .1, .01], [2, 1.5, 1.0, 0.5], delta, N)

bernoulli_default = Problem("Bernoulli", "Bernoulli",
                            [1, .1, .01], [0.9, 0.6, 0.3], delta, N)

poisson_default = Problem("Poisson", "Bernoulli",
                          [1, .1, .01], [1.5, 1.0, 0.5], delta, N)

gaussian_rev = Problem("Gaussian", "Bernoulli",
                           [.01, .1, 1], [1.5, 1.0, 0.5], delta, N)

bernoulli_rev = Problem("Bernoulli", "Bernoulli",
                            [.01, .1, 1], [0.9, 0.6, 0.3], delta, N)

poisson_rev = Problem("Poisson", "Bernoulli",
                          [.01, .1, 1], [1.5, 1.0, 0.5], delta, N)

deterministic = Problem("Deterministic", "Deterministic",
                        [1, .1, .01], [1.5, 1.0, 0.5], delta, N)

Problems = []
Problems.append(gaussian_default)
Problems.append(bernoulli_default)
Problems.append(poisson_default)
Problems.append(gaussian_rev)
Problems.append(bernoulli_rev)
Problems.append(poisson_rev)
