#!/usr/bin/env python3

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

# ------------------------------------------------------------
# Geometry
# ------------------------------------------------------------

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(lon):
    return latlon_to_vec(0.0, lon + 90.0)

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,
        dtype={"EvtDate": str}
    )

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

def fit_two_plane_model(df, label, depth_weighted=False):

    print("\n====================================================")
    print(f"TWO-PLANE FIT: {label}")
    print("Depth weighted:", depth_weighted)
    print("====================================================")

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

    y = df["abs_dT"].values

    if depth_weighted:
        d = df["TurnDep"].values
        w = (d - d.min()) / (d.max() - d.min() + 1e-9)
    else:
        w = np.ones_like(y)

    best_lon = None
    best_var = 1e9
    best_coeff = None

    # search all equatorial orientations
    for lon in np.arange(0, 180, 0.5):

        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

        # weighted least squares
        W = np.diag(w)
        Aw = W @ A
        yw = W @ y

        coeff, _, _, _ = np.linalg.lstsq(Aw, yw, rcond=None)

        y_pred = A @ coeff
        residuals = y - y_pred

        var = np.average(residuals**2, weights=w)

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

    print("Best plane 1:", best_lon, "°")
    print("Best plane 2:", (best_lon + 90) % 180, "°")
    print("Coefficients [a, b, c]:", best_coeff)
    print("Residual variance:", best_var)

    return best_lon

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

def main():

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

    df = load_data(sys.argv[1])

    subsets = {
        "ξ ≤ 15°": df[df["Ksi"] <= 15],
        "ξ ≤ 35°": df[df["Ksi"] <= 35],
        "ALL DATA": df,
        "TurnDep > 6000": df[df["TurnDep"] > 6000],
        "Deep + ξ ≤ 35°": df[(df["TurnDep"] > 6000) & (df["Ksi"] <= 35)]
    }

    for name, sub in subsets.items():

        if len(sub) < 30:
            continue

        fit_two_plane_model(sub, name, depth_weighted=False)
        fit_two_plane_model(sub, name, depth_weighted=True)

if __name__ == "__main__":
    main()
