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

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

T_MAX = 250_000
DT = 50
STEPS = int(T_MAX / DT)

TAU_I = 500_000
SIGMA_I = 0.08

Xc = 2.5
Yc = 2.5

MIN_EVENT_GAP = 10_000

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

geomagia = pd.read_csv("geomagia_combined.csv")
windows = pd.read_csv("windows.csv", sep=";")

geomagia = geomagia.sort_values("Age_ka")

# ------------------------------------------------
# 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
# ------------------------------------------------

def compute_vgp_velocity(df):

    ages = df["Age_ka"].values
    lat = df["VGP_lat"].values
    lon = df["VGP_lon"].values

    velocities = []
    times = []

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

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

        if dt == 0:
            continue

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

        v = dist/dt

        velocities.append(v)
        times.append(ages[i])

    return np.array(times), np.array(velocities)


times_ka, vgp_vel = compute_vgp_velocity(geomagia)

# normalize forcing

vgp_vel = vgp_vel / np.std(vgp_vel)

# ------------------------------------------------
# INTERPOLATE TO MODEL TIMESTEP
# ------------------------------------------------

vgp_sigma = np.std(vgp_vel)

def geomagnetic_forcing():
    return np.random.normal(0, vgp_sigma)

# ------------------------------------------------
# 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()

# ------------------------------------------------
# MANTLE FORCING
# ------------------------------------------------

def mantle_forcing():

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

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

    return base

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

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 = geomagnetic_forcing[i]

        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 WITH EXCURSION WINDOWS
# ------------------------------------------------

def compare_windows(events):

    matches=[]

    for e in events:

        age_ka = e/1000

        for _,row in windows.iterrows():

            if row.Min_age_ka <= age_ka <= row.Max_age_ka:

                matches.append((age_ka,row.Label))

    return matches

# ------------------------------------------------
# PLOTS
# ------------------------------------------------

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

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


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

    plt.show()

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

xs,ys,events = run_sim()

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

phase_plot(xs,ys)
density(xs,ys)

matches = compare_windows(events)

print("\nMatches with Frost windows:")

for m in matches:
    print(m)