import matplotlib.pyplot as plt
import numpy as np

from matplotlib.projections import PolarAxes
from matplotlib.transforms import Affine2D
from mpl_toolkits.axisartist import GridHelperCurveLinear, HostAxes


def polar_plot(fig):
    # Adapted from https://matplotlib.org/stable/gallery/axisartist/demo_floating_axis.html
    tr = (
        Affine2D().scale(np.pi / 180., 1.) +
        PolarAxes.PolarTransform(apply_theta_transforms=False)
    )
    grid_helper = GridHelperCurveLinear(tr)
    ax1 = fig.add_subplot(axes_class=HostAxes, grid_helper=grid_helper)
    ax1.axis["lon"] = ax1.new_floating_axis(1, 9)

    ax1.set_xlim(-5, 12)
    ax1.set_ylim(-5, 10)
    ax1.grid(True)
    for key, val in ax1.axis.items():
        val.major_ticklabels.set_visible(False)


def main():
    for layout in [None, "tight", "constrained"]:
        my_fig = plt.figure(figsize=(5, 5), layout=layout)
        polar_plot(my_fig)
        layout_str = layout if layout is not None else "default"
        plt.title(f"Layout: {layout_str}")
        plt.savefig(f"floating_axis_{layout_str}.png")


if __name__ == "__main__":
    main()