#!/usr/bin/env python3

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

EARTH_RADIUS = 6371.0

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

def sph_to_cart(lat, lon, r):
    lat = radians(lat)
    lon = radians(lon)
    x = r * cos(lat) * cos(lon)
    y = r * cos(lat) * sin(lon)
    z = r * sin(lat)
    return x, y, z

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 Frost 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

# ------------------------------------------------------------
# Write VTK
# ------------------------------------------------------------

def write_vtk(filename, xyz, scalars):

    n = len(xyz)

    with open(filename, "w") as f:
        f.write("# vtk DataFile Version 3.0\n")
        f.write("Frost inner core turning points\n")
        f.write("ASCII\n")
        f.write("DATASET POLYDATA\n")

        f.write(f"POINTS {n} float\n")
        for x,y,z in xyz:
            f.write(f"{x} {y} {z}\n")

        f.write(f"\nVERTICES {n} {n*2}\n")
        for i in range(n):
            f.write(f"1 {i}\n")

        f.write(f"\nPOINT_DATA {n}\n")

        for name, values in scalars.items():
            f.write(f"\nSCALARS {name} float 1\n")
            f.write("LOOKUP_TABLE default\n")
            for v in values:
                f.write(f"{v}\n")

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

def main():

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

    df = load_data(sys.argv[1])

    best_lon, coeff = fit_two_plane(df)
    print("Best planes:", best_lon, "and", best_lon+90)

    # 3D positions
    xyz = []
    unit_vecs = []

    for la, lo, dep in zip(df["TurnLa"], df["TurnLo"], df["TurnDep"]):
        r = EARTH_RADIUS - dep
        x,y,z = sph_to_cart(la, lo, r)
        xyz.append((x,y,z))
        unit_vecs.append(latlon_to_vec(la, lo))

    unit_vecs = np.array(unit_vecs)

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

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

    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"].values - pred

    scalars = {
        "abs_dT": df["abs_dT"].values,
        "residual": residuals,
        "depth": df["TurnDep"].values,
        "dist_plane1": d1,
        "dist_plane2": d2
    }

    write_vtk("frost_innercore.vtk", xyz, scalars)

    print("Wrote frost_innercore.vtk")

if __name__ == "__main__":
    main()
