import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter

# ------------------------------------------------
# RANDOM SEED (optional reproducibility)
# ------------------------------------------------

np.random.seed(1)

# ------------------------------------------------
# MODEL PARAMETERS
# ------------------------------------------------

T_MAX = 1_000_000        # years simulated
DT = 50                  # timestep

STEPS = int(T_MAX / DT)

# relaxation times
TAU_I = 500_000          # mantle inertia relaxation
TAU_B = 20_000           # magnetic torque relaxation

# noise amplitudes
SIGMA_I = 0.08
SIGMA_B = 0.20

# MTDE instability ellipse
Xc = 2.5
Yc = 2.5

# refractory period between events
MIN_EVENT_GAP = 10_000   # years


# ------------------------------------------------
# STOCHASTIC FORCING
# ------------------------------------------------

def mantle_forcing():

    base = np.random.normal(0,1)

    # rare plume pulse
    if np.random.rand() < 3e-5:
        return base + np.random.normal(6,1)

    return base


def geomagnetic_forcing():

    return np.random.normal(0,1)


# ------------------------------------------------
# ORNSTEIN-UHLENBECK STEP
# ------------------------------------------------

def ou_step(x, tau, sigma):

    a = np.exp(-DT/tau)

    return a*x + sigma*np.sqrt(1-a*a)*np.random.normal()


# ------------------------------------------------
# SIMULATION
# ------------------------------------------------

def run_sim():

    X = 0
    Y = 0

    xs=[]
    ys=[]
    events=[]

    prev_phi = 0
    last_event = -1e12

    for i in range(STEPS):

        # OU dynamics + forcing
        X = ou_step(X,TAU_I,SIGMA_I) + 0.02*mantle_forcing()
        Y = ou_step(Y,TAU_B,SIGMA_B) + 0.05*geomagnetic_forcing()

        xs.append(X)
        ys.append(Y)

        phi = (X/Xc)**2 + (Y/Yc)**2

        t = i*DT

        if phi>1 and prev_phi<=1 and (t-last_event)>MIN_EVENT_GAP:

            events.append(t)
            last_event = t

        prev_phi = phi

    return np.array(xs), np.array(ys), np.array(events)


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

def phase_plot(xs,ys):

    plt.figure(figsize=(7,7))

    plt.scatter(xs,ys,s=1,alpha=0.1)

    t = np.linspace(0,2*np.pi,400)

    plt.plot(Xc*np.cos(t),Yc*np.sin(t),'r',lw=2)

    plt.xlim(-3,3)
    plt.ylim(-3,3)

    plt.xlabel("dI/dt (Inertia Forcing)")
    plt.ylabel("dB/dt (Magnetic Torque)")

    plt.title("MTDE Torque–Inertia Phase Space")

    plt.show()


# ------------------------------------------------
# PROBABILITY DENSITY
# ------------------------------------------------

def density(xs,ys):

    H,x,y = np.histogram2d(xs,ys,bins=200)

    H = gaussian_filter(H,2)

    plt.figure(figsize=(7,7))

    plt.imshow(
        H.T,
        origin='lower',
        cmap='inferno',
        extent=[x[0],x[-1],y[0],y[-1]]
    )

    t = np.linspace(0,2*np.pi,400)

    plt.plot(Xc*np.cos(t),Yc*np.sin(t),'cyan')

    plt.xlim(-3,3)
    plt.ylim(-3,3)

    plt.xlabel("dI/dt")
    plt.ylabel("dB/dt")

    plt.title("MTDE Probability Density")

    plt.colorbar(label="Probability Density")

    plt.show()


# ------------------------------------------------
# EVENT STATISTICS
# ------------------------------------------------

def stats(events):

    if len(events)<2:

        print("Too few events")
        return

    intervals = np.diff(events)

    print("\nMTDE EVENT STATS")
    print("----------------")

    print("Events:",len(events))
    print("Mean recurrence:",np.mean(intervals)/1000,"kyr")
    print("Median:",np.median(intervals)/1000,"kyr")
    print("Std:",np.std(intervals)/1000,"kyr")

    plt.figure()

    plt.hist(intervals/1000,bins=15)

    plt.xlabel("Recurrence Interval (kyr)")
    plt.ylabel("Count")

    plt.title("MTDE Recurrence")

    plt.show()


# ------------------------------------------------
# MAIN
# ------------------------------------------------

xs,ys,events = run_sim()

print("Total MTDE events:",len(events))

phase_plot(xs,ys)

density(xs,ys)

stats(events)