#!/usr/bin/env python3

import sys
import numpy as np
import pandas as pd
from math import radians, sin, cos, asin

# ------------------------------------------------------------
# Geometry helpers
# ------------------------------------------------------------

def latlon_to_vec(lat, lon):
    lat = radians(lat)
    lon = radians(lon)
    x = cos(lat)*cos(lon)
    y = cos(lat)*sin(lon)
    z = sin(lat)
    return np.array([x,y,z])

def plane_normal_from_equatorial_lon(lon):
    return latlon_to_vec(0.0, lon + 90.0)

def angular_distance_to_plane(point_vec, plane_normal):
    dot = abs(np.dot(point_vec, plane_normal))
    dot = min(1.0, max(-1.0, dot))
    return np.degrees(asin(dot))

# ------------------------------------------------------------
# Load dataset
# ------------------------------------------------------------

def load_frost_dat(path):

    with open(path) as f:
        lines = f.readlines()

    header_idx = None
    for i, line in enumerate(lines):
        if line.strip().startswith("EvtDate"):
            header_idx = i
            break

    df = pd.read_csv(
        path,
        sep=r"\s+",
        engine="python",
        skiprows=header_idx,
        dtype={"EvtDate": str, "Sta": str, "DataFromSource": str}
    )

    # unified signal
    df["dT"] = df["dTabdf"].copy()
    mask = df["dT"] == 0
    df.loc[mask, "dT"] = df.loc[mask, "dTbcdf"]

    df["abs_dT"] = df["dT"].abs()

    return df

# ------------------------------------------------------------
# Plane sweep
# ------------------------------------------------------------

def full_geometry_sweep(df, label):

    print("\n====================================================")
    print(f" GEOMETRY SWEEP: {label}")
    print("====================================================")

    points = np.array([
        latlon_to_vec(la, lo)
        for la, lo in zip(df["TurnLa"], df["TurnLo"])
    ])

    signed = df["dT"].values
    strength = df["abs_dT"].values

    results = []

    for lon in np.arange(0, 180, 1):

        n = plane_normal_from_equatorial_lon(lon)

        dists = np.array([
            angular_distance_to_plane(p, n)
            for p in points
        ])

        corr_strength = np.corrcoef(dists, strength)[0,1]
        corr_signed   = np.corrcoef(dists, signed)[0,1]

        near = strength[dists < 10]
        far  = strength[dists > 30]

        if len(near) > 5 and len(far) > 5:
            contrast = near.mean() - far.mean()
        else:
            contrast = np.nan

        results.append((lon, corr_strength, corr_signed, contrast))

    res = pd.DataFrame(results, columns=[
        "lon",
        "corr_strength",
        "corr_signed",
        "contrast"
    ])

    best_pos = res.loc[res["corr_strength"].idxmin()]
    best_neg = res.loc[res["corr_strength"].idxmax()]
    best_signed = res.loc[res["corr_signed"].abs().idxmax()]

    print("\n--- Strongest anisotropy concentration plane ---")
    print(best_pos)

    print("\n--- Strongest depletion plane ---")
    print(best_neg)

    print("\n--- Strongest signed-structure plane ---")
    print(best_signed)

    # Orthogonal pairing
    print("\n--- Orthogonal partner ---")
    lon = best_signed["lon"]
    ortho = (lon + 90) % 180
    match = res.iloc[(res["lon"] - ortho).abs().argsort()[:1]]
    print(match.iloc[0].to_dict())

    # Direct tests
    test_specific_axis(df, -60)
    test_specific_axis(df, 120)

    return res

# ------------------------------------------------------------
# Specific plane tests
# ------------------------------------------------------------

def test_specific_axis(df, lon):

    n = plane_normal_from_equatorial_lon(lon)

    points = np.array([
        latlon_to_vec(la, lo)
        for la, lo in zip(df["TurnLa"], df["TurnLo"])
    ])

    dists = np.array([
        angular_distance_to_plane(p, n)
        for p in points
    ])

    print(f"\nPlane test @ {lon}°:")
    print("corr(|dT|):", np.corrcoef(dists, df["abs_dT"])[0,1])
    print("corr(dT):  ", np.corrcoef(dists, df["dT"])[0,1])

# ------------------------------------------------------------
# Main driver
# ------------------------------------------------------------

def main():

    if len(sys.argv) != 2:
        print("Usage: python frost_analysis_full.py <datafile>")
        sys.exit(1)

    df = load_frost_dat(sys.argv[1])

    print("\nLoaded rows:", len(df))

    # --- Regime 1: polar (baseline)
    polar15 = df[df["Ksi"] <= 15]
    full_geometry_sweep(polar15, "ξ ≤ 15°")

    # --- Regime 2: extended polar
    polar35 = df[df["Ksi"] <= 35]
    full_geometry_sweep(polar35, "ξ ≤ 35°")

    # --- Regime 3: all data
    full_geometry_sweep(df, "ALL DATA")

    # --- Regime 4: deep inner core
    deep = df[df["TurnDep"] > 6000]
    full_geometry_sweep(deep, "TurnDep > 6000 km")

    # --- Regime 5: deep + polar
    deep_polar = deep[deep["Ksi"] <= 35]
    full_geometry_sweep(deep_polar, "TurnDep > 6000 km AND ξ ≤ 35°")


if __name__ == "__main__":
    main()
