import numpy as np
import copy
from scipy.optimize import linprog, minimize
from tabulate import tabulate
from tqdm import tqdm
import matplotlib.pyplot as plt

np.random.seed(457)


class FarmingProblem:
    def __init__(self):
        super().__init__()
        self.c = np.array([150, 230, 260, 238, 210, -170, -150, -36, -10])
        self.A = np.array([[1, 1, 1, 0, 0, 0, 0, 0, 0],
                           [-2.5, 0, 0, -1, 0, 1, 0, 0, 0],
                           [0, -3, 0, 0, -1, 0, 1, 0, 0],
                           [0, 0, -20, 0, 0, 0, 0, 1, 1],
                           [0, 0, 0, 0, 0, 0, 0, 1, 0]])
        self.b = np.array([500, -200, -240, 0, 6000])

        self.stds = [50, 40, 16, 5]

        solution = linprog(self.c, A_ub=self.A, b_ub=self.b)
        self.optimal_val = solution['fun']
        self.optimal_x = solution['x']

    def __call__(self, x):
        temp_c = copy.deepcopy(self.c)
        temp_c[5] = np.random.normal(self.c[5], self.stds[0])
        temp_c[6] = np.random.normal(self.c[6], self.stds[1])
        temp_c[7] = np.random.normal(self.c[7], self.stds[2])
        return temp_c @ x.T

    def expected(self, x):
        return self.c @ x.T

    def estimate_gradient(self, m, return_var=False):
        # Estimate the gradient at x
        # The gradient is the same in every point as the objective function is linear, so just sample random
        # coefficients for the function m times and use it as gradient estimate
        gradient_estimates = []
        for _ in range(m):
            temp_c = copy.deepcopy(self.c)
            temp_c[5] = np.random.normal(self.c[5], self.stds[0])
            temp_c[6] = np.random.normal(self.c[6], self.stds[1])
            temp_c[7] = np.random.normal(self.c[7], self.stds[2])
            gradient_estimates.append(temp_c)
        if return_var:
            return np.mean(gradient_estimates, axis=0), self.stds
        else:
            return np.mean(gradient_estimates, axis=0)

    def projection(self, x):
        # Calculate the orthogonal projection back on the convex set of constraints
        loss = lambda x_: np.linalg.norm((x - x_))
        cons = {'type': 'ineq',
                'fun': lambda x_: self.b - np.dot(self.A, x_),
                'jac': lambda x_: -self.A}
        bounds = [(0, None) for _ in range(x.shape[0])]

        x_proj = minimize(loss, x, jac='cs', constraints=cons, bounds=bounds, method='SLSQP', options={'disp': False})['x']
        return x_proj


def rspg_algorithm(x_0, gammas, f, S, m):
    candidate_solutions = np.zeros((S, x_0.shape[0]))
    run_losses = np.zeros(S)
    all_losses, all_R = [], []
    # Run algorithm S times for random R
    for i in range(S):
        losses = []
        R = np.random.randint(1, len(gammas))
        all_R.append(R)
        x = copy.deepcopy(x_0)
        # Do step of RSPG algorithm for each gamma_k, k = 1, ..., R
        for gamma in tqdm(gammas[:R]):
            y = x - gamma * f.estimate_gradient(m)
            x = f.projection(y)
            loss = np.sqrt((f.expected(x) - f.optimal_val) ** 2)
            losses.append(loss)
        print(f'\tLoss: {loss}')
        all_losses.append(losses)
        candidate_solutions[i] = x
        run_losses[i] = loss
    # Select the run with lowest error as optimal solution
    best_loss = np.argmin(run_losses)
    return candidate_solutions[best_loss], all_losses, all_R


if __name__ == '__main__':
    func = FarmingProblem()
    data_shape = 9

    M = 500

    sigma = np.sqrt(np.sum(np.array(func.stds) ** 2))
    alpha = 1
    L = 0.05

    # Start at some initial plausible state
    x_0 = np.array([0, 0, 0, -200, -240, 0, 0, 0, 0])
    D = 6000

    m = int(np.ceil(min([M, max([1, (sigma * np.sqrt(6 * M)) / (4 * L * D)])])))

    x_opt, losses, Rs = rspg_algorithm(
        x_0=x_0,
        gammas=[1 / (2 * L)] * int(np.floor(M / m)),
        f=func,
        S=6,
        m=m
    )
    print()
    print(f'===============')
    print(f'Wheat planted:\t\t {x_opt[0]:.0f} ar\n'
          f'Corn planted:\t\t {x_opt[1]:.0f} ar\n'
          f'Beats planted:\t\t {x_opt[2]:.0f} ar\n')
    print('-----------------------------------------------------------------------------------------------------')
    print(tabulate([[f'{x_opt[3]:.0f}', f'{x_opt[4]:.0f}', f'{x_opt[5]:.0f}', f'{x_opt[6]:.0f}', f'{x_opt[7]:.0f}']],
                   headers=['Wheat bought', 'Corn bought', 'Wheat sold', 'Corn sold', 'Beats sold']))
    print(f'Expected Profit:\t{-func.expected(x_opt):.0f}$')

    fig, axes = plt.subplots(2, 3, figsize=(8, 4), sharey=True, sharex=True)
    colors = ['red', 'blue', 'green', 'orange', 'purple', 'cyan', 'black', 'brown']
    Rs = sorted(Rs, reverse=True)
    losses = sorted(losses, key=lambda x: len(x), reverse=True)
    for i, (R, loss, ax) in enumerate(zip(Rs, losses, axes.flat)):
        last_loss = loss[-1]
        idx = len(loss) - 1
        ax.plot(loss, c='black')
        ax.scatter(x=idx, y=loss[-1], c='black')
        ax.set_title(f'R={R},\n Chyba={last_loss:.0f}')
    fig.supxlabel('Počet kroků')
    fig.supylabel(r"Chyba ($\psi(x_R) - \psi^*$)")
    plt.tight_layout()
    plt.savefig('loss.pdf')