#!/usr/bin/env python3

import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from math import radians, sin, cos, asin

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

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

def plane_normal(lon):
    return latlon_to_vec(0.0, lon + 90)

def ang_dist_to_plane(p, n):
    dot = abs(np.dot(p, n))
    dot = min(1.0, max(-1.0, dot))
    return np.degrees(asin(dot))

# ------------------------------------------------------------
# Load data
# ------------------------------------------------------------

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

    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

# ------------------------------------------------------------
# Two-plane fit
# ------------------------------------------------------------

def fit_two_plane(df):

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

    y = df["abs_dT"].values

    best_lon = None
    best_var = 1e9

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

        n1 = plane_normal(lon)
        n2 = plane_normal(lon + 90)

        d1 = np.array([ang_dist_to_plane(p, n1) for p in points])
        d2 = np.array([ang_dist_to_plane(p, n2) for p in points])

        x1 = np.cos(np.radians(d1))**2
        x2 = np.cos(np.radians(d2))**2

        A = np.vstack([np.ones_like(x1), x1, x2]).T
        coeff, _, _, _ = np.linalg.lstsq(A, y, rcond=None)

        y_pred = A @ coeff
        var = np.var(y - y_pred)

        if var < best_var:
            best_var = var
            best_lon = lon
            best_coeff = coeff

    return best_lon, best_coeff

# ------------------------------------------------------------
# Visualization suite
# ------------------------------------------------------------

def plot_turning_points(df):
    plt.figure(figsize=(10,5))
    plt.scatter(df["TurnLo"], df["TurnLa"],
                c=df["abs_dT"], s=8)
    plt.colorbar(label="|dT|")
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")
    plt.title("Inner Core Turning Points Colored by |dT|")
    plt.savefig("01_turning_points_map.png", dpi=200)
    plt.close()

def plot_longitude_response(df):
    plt.figure(figsize=(10,5))
    plt.scatter(df["TurnLo"], df["abs_dT"], s=6)
    plt.xlabel("Turn Longitude")
    plt.ylabel("|dT|")
    plt.title("|dT| vs Longitude")
    plt.savefig("02_longitude_vs_dT.png", dpi=200)
    plt.close()

def plot_depth_structure(df):
    plt.figure(figsize=(10,5))
    plt.scatter(df["TurnDep"], df["abs_dT"], s=6)
    plt.xlabel("Turning Depth (km)")
    plt.ylabel("|dT|")
    plt.title("|dT| vs Depth")
    plt.savefig("03_depth_vs_dT.png", dpi=200)
    plt.close()

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

    n1 = plane_normal(best_lon)
    dists = np.array([ang_dist_to_plane(p, n1) for p in points])

    plt.figure(figsize=(10,5))
    plt.scatter(dists, df["abs_dT"], s=6)
    plt.xlabel("Distance to Best Plane (deg)")
    plt.ylabel("|dT|")
    plt.title("Anisotropy vs Distance to Best Plane")
    plt.savefig("04_distance_response.png", dpi=200)
    plt.close()

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

    n1 = plane_normal(best_lon)
    n2 = plane_normal(best_lon + 90)

    d1 = np.array([ang_dist_to_plane(p, n1) for p in points])
    d2 = np.array([ang_dist_to_plane(p, n2) for p in points])

    x1 = np.cos(np.radians(d1))**2
    x2 = np.cos(np.radians(d2))**2

    pred = coeff[0] + coeff[1]*x1 + coeff[2]*x2
    residuals = df["abs_dT"] - pred

    plt.figure(figsize=(10,5))
    plt.scatter(df["TurnLo"], df["TurnLa"],
                c=residuals, s=8)
    plt.colorbar(label="Residual")
    plt.title("Two-Plane Model Residuals")
    plt.savefig("05_residual_map.png", dpi=200)
    plt.close()

def plot_polar_histogram(best_lon):
    angles = np.radians([best_lon, (best_lon+90)%180])
    plt.figure(figsize=(6,6))
    ax = plt.subplot(111, polar=True)
    ax.bar(angles, [1,1], width=0.2)
    ax.set_title("Best-Fit Plane Orientations")
    plt.savefig("06_plane_orientation_polar.png", dpi=200)
    plt.close()

# ------------------------------------------------------------
# Main
# ------------------------------------------------------------

def main():

    if len(sys.argv) != 2:
        print("Usage: python frost_visualize.py <datafile>")
        return

    df = load_data(sys.argv[1])

    best_lon, coeff = fit_two_plane(df)

    print("Best plane:", best_lon)
    print("Orthogonal:", (best_lon+90)%180)

    plot_turning_points(df)
    plot_longitude_response(df)
    plot_depth_structure(df)
    plot_plane_distance(df, best_lon)
    plot_residual_map(df, best_lon, coeff)
    plot_polar_histogram(best_lon)

if __name__ == "__main__":
    main()
