"""Kirsch boundary stress for a circular hole under remote horizontal tension."""
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import Normalize
import numpy as np


def main():
    theta = np.linspace(0, 2*np.pi, 721)
    tangential = 1 - 2*np.cos(2*theta)
    assert np.isclose(tangential.min(), -1)
    assert np.isclose(tangential.max(), 3)
    assert np.allclose(1 - 2*np.cos(2*np.array([0, np.pi/2, np.pi, 3*np.pi/2])),
                       [-1, 3, -1, 3])
    fig, (hole, curve) = plt.subplots(1, 2, figsize=(9, 4.1), dpi=100,
                                     layout="constrained", width_ratios=[1, 1.15])
    fig.set_facecolor("white")
    points = np.column_stack([np.cos(theta), np.sin(theta)])
    lines = LineCollection(np.stack([points[:-1], points[1:]], axis=1),
                           cmap="coolwarm", norm=Normalize(-1, 3), linewidth=6)
    lines.set_array((tangential[:-1]+tangential[1:])/2)
    hole.add_collection(lines)
    for y in [-.7, 0, .7]:
        for start, end in [(1.7, 2.5), (-1.7, -2.5)]:
            hole.annotate("", (end, y), (start, y),
                          arrowprops=dict(arrowstyle="->", color="#555555", lw=1.5))
    hole.text(0, 1.33, "3σ∞", ha="center", color="#aa2222", fontsize=12)
    hole.text(0, -1.55, "3σ∞", ha="center", color="#aa2222", fontsize=12)
    hole.text(0, 0, "Hole\n(radial traction = 0)", ha="center", va="center", fontsize=10)
    hole.text(0, 1.9, "Remote tension along x", ha="center")
    hole.set(xlim=(-2.8, 2.8), ylim=(-2, 2.3), aspect="equal",
             title="Tangential stress at the boundary")
    hole.axis("off")
    curve.plot(np.degrees(theta), tangential, color="#1565c0", lw=2)
    curve.axhline(0, color="#aaaaaa", lw=.8)
    curve.scatter([0, 90, 180, 270, 360], [-1, 3, -1, 3, -1],
                  color=["#1565c0", "#aa2222", "#1565c0", "#aa2222", "#1565c0"])
    curve.set(xlim=(0, 360), ylim=(-1.5, 3.5), xticks=[0, 90, 180, 270, 360],
              yticks=[-1, 0, 1, 2, 3], xlabel="Angle θ from loading direction (degrees)",
              ylabel="Tangential stress / remote tensile stress",
              title="σθθ(a, θ) / σ∞ = 1 − 2 cos(2θ)")
    curve.grid(alpha=.2)
    fig.savefig(Path(__file__).with_suffix(".png").name, facecolor="white")
    plt.close(fig)


if __name__ == "__main__":
    main()