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

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

T_MAX = 250_000      # simulate last 250 kyr
DT = 50

STEPS = int(T_MAX / DT)

TAU_I = 500_000
TAU_B = 20_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=";")


# ------------------------------------------------
# MAGNETIC NOISE FROM GEOMAGIA
# ------------------------------------------------

def compute_geomagnetic_sigma(df):

    # convert VGP scatter to angular variance
    lat = np.radians(df["VGP_lat"])
    lon = np.radians(df["VGP_lon"])

    # simple angular dispersion
    sigma_lat = np.std(lat)
    sigma_lon = np.std(lon)

    sigma = np.sqrt(sigma_lat**2 + sigma_lon**2)

    return sigma


SIGMA_B = compute_geomagnetic_sigma(geomagia)


# ------------------------------------------------
# FORCING FUNCTIONS
# ------------------------------------------------

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)


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


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


# ------------------------------------------------
# COMPARE WITH EXCURSION WINDOWS
# ------------------------------------------------

def compare_windows(events,windows):

    matches=[]

    for e in events:

        for _,row in windows.iterrows():

            if row.Min_age_ka*1000 <= e <= row.Max_age_ka*1000:

                matches.append((e,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',lw=2)

    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(label="Probability Density")

    plt.show()


def stats(events):

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

matches = compare_windows(events,windows)

print("\nEvents matching Frost (2020) windows:")

for m in matches:
    print(m)