#!/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}
    )

    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

# ------------------------------------------------------------
# Depth weighting
# ------------------------------------------------------------

def depth_weight(df):
    """
    Normalize depth so deeper rays get more influence.
    """
    d = df["TurnDep"].values
    return (d - d.min()) / (d.max() - d.min() + 1e-9)

# ------------------------------------------------------------
# Depth-weighted sweep
# ------------------------------------------------------------

def depth_weighted_sweep(df, label):

    print("\n====================================================")
    print(f" DEPTH-WEIGHTED SWEEP: {label}")
    print("====================================================")

    weights = depth_weight(df)

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

    strength = df["abs_dT"].values

    best_lon = None
    best_score = 1e9

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

        model = np.cos(np.radians(dists))**2

        # weighted least squares residual
        residual = np.average((strength - model)**2, weights=weights)

        if residual < best_score:
            best_score = residual
            best_lon = lon

    print("Depth-weighted best-fit plane:", best_lon, "°")
    print("Residual score:", best_score)

    return best_lon

# ------------------------------------------------------------
# Plane optimization (residual minimization)
# ------------------------------------------------------------

def optimize_plane_residuals(df, label):

    print("\n====================================================")
    print(f" RESIDUAL MINIMIZATION: {label}")
    print("====================================================")

    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, 0.5):

        n = plane_normal_from_equatorial_lon(lon)

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

        x = np.cos(np.radians(dists))**2

        # linear fit y = a + b*x
        A = np.vstack([np.ones_like(x), x]).T
        coeff, _, _, _ = np.linalg.lstsq(A, y, rcond=None)

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

        if var < best_var:
            best_var = var
            best_lon = lon

    print("Plane minimizing anisotropy residuals:", best_lon, "°")
    print("Residual variance:", best_var)

    return best_lon

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

def main():

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

    df = load_frost_dat(sys.argv[1])

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

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

        depth_weighted_sweep(sub, name)
        optimize_plane_residuals(sub, name)

if __name__ == "__main__":
    main()
