import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# ------------------------------------------------
# PARAMETERS
# ------------------------------------------------

T_MAX = 250000
DT = 50
STEPS = int(T_MAX / DT)

TAU_I = 500000
SIGMA_I = 0.08

Xc = 2.5
Yc = 2.5

MIN_EVENT_GAP = 10000

TIME_BIN_YR = 1000


# ------------------------------------------------
# LOAD GEOMAGIA DATA
# ------------------------------------------------

df = pd.read_csv("geomagia_combined.csv")

df = df.sort_values("Age_ka")


# ------------------------------------------------
# TIME BINNING (remove dt→0 artifacts)
# ------------------------------------------------

df["age_yr"] = df["Age_ka"] * 1000
df["age_bin"] = (df["age_yr"] // TIME_BIN_YR) * TIME_BIN_YR

binned = df.groupby("age_bin").mean(numeric_only=True)

ages = binned.index.values / 1000
lat = binned["VGP_lat"].values
lon = binned["VGP_lon"].values


# ------------------------------------------------
# GREAT CIRCLE DISTANCE
# ------------------------------------------------

def angular_distance(lat1, lon1, lat2, lon2):

    lat1 = np.radians(lat1)
    lon1 = np.radians(lon1)

    lat2 = np.radians(lat2)
    lon2 = np.radians(lon2)

    d = np.sin(lat1)*np.sin(lat2) + np.cos(lat1)*np.cos(lat2)*np.cos(lon2-lon1)

    return np.degrees(np.arccos(np.clip(d,-1,1)))


# ------------------------------------------------
# COMPUTE VGP ANGULAR VELOCITY
# ------------------------------------------------

vel = []
vel_age = []

for i in range(1,len(ages)):

    dt = abs(ages[i] - ages[i-1])

    if dt == 0:
        continue

    d = angular_distance(lat[i-1], lon[i-1], lat[i], lon[i])

    vel.append(d/dt)
    vel_age.append(ages[i])

vel = np.array(vel)
vel_age = np.array(vel_age)


# ------------------------------------------------
# EXCURSION DETECTION
# ------------------------------------------------

threshold = np.mean(vel) + 2*np.std(vel)

excursions = []
current = None

for i,v in enumerate(vel):

    if v > threshold:

        if current is None:
            current = [vel_age[i], vel_age[i]]

        current[1] = vel_age[i]

    else:

        if current is not None:
            excursions.append(tuple(current))
            current = None

if current is not None:
    excursions.append(tuple(current))


print("\nDetected excursion candidates:")
for e in excursions:
    print(e)


# ------------------------------------------------
# OU STEP
# ------------------------------------------------

def ou_step(x, tau, sigma):

    a = np.exp(-DT/tau)

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


# ------------------------------------------------
# FORCING
# ------------------------------------------------

sigma_mag = np.std(vel)

def magnetic_forcing():
    return np.random.normal(0, sigma_mag)


def mantle_forcing():

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

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

    return base


# ------------------------------------------------
# RUN MTDE MODEL
# ------------------------------------------------

def run_sim():

    X = 0
    Y = 0

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

    prev_phi = 0
    last_event = -1e12

    for i in range(STEPS):

        X = ou_step(X, TAU_I, SIGMA_I) + 0.02 * mantle_forcing()
        Y = magnetic_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)


# ------------------------------------------------
# COMPARE EVENTS TO EXCURSIONS
# ------------------------------------------------

def compare(events):

    matches = []

    for e in events:

        age = e / 1000

        for start,end in excursions:

            if start <= age <= end:
                matches.append((age,start,end))

    return matches


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

xs,ys,events = run_sim()

print("\nMTDE events:",len(events))

matches = compare(events)

print("\nMTDE events overlapping excursions:")
for m in matches:
    print(m)


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

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')

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

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

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

plt.show()


# ------------------------------------------------
# VELOCITY DIAGNOSTIC PLOT
# ------------------------------------------------

plt.figure()

plt.plot(vel_age,vel)

plt.axhline(threshold,color='r')

plt.xlabel("Age (ka)")
plt.ylabel("VGP angular velocity (deg/kyr)")

plt.title("Geomagnetic Angular Velocity")

plt.show()