from matplotlib import pyplot as plt
from leap_ec.real_rep.problems import CosineFamilyProblem, MatrixTransformedProblem, plot_2d_problem

original_problem = CosineFamilyProblem(alpha=1.0, global_optima_counts=[2, 3], local_optima_counts=[2, 3])

transformed_problem = MatrixTransformedProblem.random_orthonormal(original_problem, 2)

fig = plt.figure(figsize=(12, 8))

plt.subplot(221, projection='3d')
bounds = original_problem.bounds
plot_2d_problem(original_problem, xlim=bounds, ylim=bounds, ax=plt.gca(), granularity=0.025)

plt.subplot(222, projection='3d')
plot_2d_problem(transformed_problem, xlim=bounds, ylim=bounds, ax=plt.gca(), granularity=0.025)

plt.subplot(223)
plot_2d_problem(original_problem, kind='contour', xlim=bounds, ylim=bounds, ax=plt.gca(), granularity=0.025)

plt.subplot(224)
plot_2d_problem(transformed_problem, kind='contour', xlim=bounds, ylim=bounds, ax=plt.gca(), granularity=0.025)