#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse, FancyArrowPatch, Rectangle

plt.rcParams.update({
    "font.family": "serif",
    "font.size": 12
})


# ------------------------------------------
# PHASE SPACE FIGURE
# ------------------------------------------

def generate_phase_space():

    fig, ax = plt.subplots(figsize=(6.5,6))

    theta = np.linspace(0, 2*np.pi, 400)
    x = np.cos(theta)
    y = np.sin(theta)

    # stability boundary
    ax.plot(x, y, linewidth=2, label="MTDE stability boundary")

    # unstable region shading
    ax.fill(x, y, alpha=0.15)

    # axes
    ax.axhline(0, linewidth=1)
    ax.axvline(0, linewidth=1)

    # example points
    ax.scatter(0.25, 0.1, s=80)
    ax.text(0.30,0.1,"Typical mantle\nreorganization", va='center')

    ax.scatter(-0.15,0.35, s=80)
    ax.text(-0.45,0.35,"Geomagnetic\nexcursion", va='center')

    ax.scatter(0.75,0.65, s=90)
    ax.text(0.45,0.75,"MTDE regime", weight="bold")

    # region labels
    ax.text(-0.9,-0.9,"Stable regime")
    ax.text(0.25,0.35,"Instability\nthreshold")

    ax.set_xlabel(r"$dI/dt$  (inertia forcing)")
    ax.set_ylabel(r"$dB/dt$  (magnetic torque forcing)", rotation=90)

    ax.set_xlim(-1.2,1.2)
    ax.set_ylim(-1.2,1.2)

    ax.set_aspect('equal')

    ax.set_title("MTDE Torque–Inertia Phase Space")

    plt.tight_layout()
    plt.savefig("mtde_phase_space.pdf")
    plt.close()



# ------------------------------------------
# SYSTEM ARCHITECTURE FIGURE
# ------------------------------------------

def generate_system_architecture():

    fig, ax = plt.subplots(figsize=(10,4))

    ax.axis('off')

    width = 0.20
    height = 0.25

    boxes = {

        "Mantle\nConvection": (0.05,0.40),

        "Mass\nRedistribution": (0.30,0.40),

        "Inertia Tensor\nEvolution": (0.55,0.40),

        "MTDE\nThreshold": (0.80,0.40)

    }

    for label,(x,y) in boxes.items():

        rect = Rectangle((x,y), width, height,
                         fill=False, linewidth=1.8)

        ax.add_patch(rect)

        ax.text(x+width/2,
                y+height/2,
                label,
                ha='center',
                va='center')


    arrows = [
        ((0.25,0.525),(0.30,0.525)),
        ((0.50,0.525),(0.55,0.525)),
        ((0.75,0.525),(0.80,0.525))
    ]

    for start,end in arrows:

        arrow = FancyArrowPatch(start,end,
                                arrowstyle='->',
                                mutation_scale=18,
                                linewidth=1.6)

        ax.add_patch(arrow)


    # geomagnetic arrow

    arrow = FancyArrowPatch((0.70,0.80),(0.90,0.63),
                            arrowstyle='->',
                            mutation_scale=18,
                            linewidth=1.6)

    ax.add_patch(arrow)

    ax.text(0.67,0.83,
            "Geomagnetic\nVariability",
            ha='center')


    ax.set_title(
        "Mantle Convection – Inertia – Magnetic Torque Coupling in MTDE Framework",
        pad=20
    )

    plt.tight_layout()

    plt.savefig("mtde_system_architecture.pdf")

    plt.close()



# ------------------------------------------

def main():

    generate_phase_space()

    generate_system_architecture()

    print("Corrected MTDE figures generated.")


if __name__ == "__main__":

    main()