#!/usr/bin/env python3

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

EARTH_RADIUS = 6371.0
INNER_CORE_RADIUS = 1220.0

# ------------------------------------------------------------
# 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 orthonormal_basis(n):
    """
    Build 2 vectors spanning the plane perpendicular to n
    """
    if abs(n[0]) < 0.9:
        ref = np.array([1,0,0])
    else:
        ref = np.array([0,1,0])

    u = np.cross(n, ref)
    u /= np.linalg.norm(u)

    v = np.cross(n, u)
    v /= np.linalg.norm(v)

    return u, v

# ------------------------------------------------------------
# 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 (same as before)
# ------------------------------------------------------------

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.abs(points @ n1)
        d2 = np.abs(points @ n2)

        x1 = d1**2
        x2 = 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

    return best_lon

# ------------------------------------------------------------
# Generate plane mesh
# ------------------------------------------------------------

def build_plane_mesh(normal, size=INNER_CORE_RADIUS, res=60):
    """
    Build a square mesh inside the inner core
    """

    u, v = orthonormal_basis(normal)

    pts = []
    faces = []

    for i in range(res):
        for j in range(res):
            a = (i/(res-1) - 0.5) * 2 * size
            b = (j/(res-1) - 0.5) * 2 * size

            p = a*u + b*v
            pts.append(p)

    # build quad faces
    for i in range(res-1):
        for j in range(res-1):
            id0 = i*res + j
            id1 = id0 + 1
            id2 = id0 + res + 1
            id3 = id0 + res

            faces.append([id0, id1, id2, id3])

    return np.array(pts), faces

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

def write_plane_vtk(filename, planes):

    all_pts = []
    all_faces = []
    offset = 0

    for pts, faces in planes:
        all_pts.extend(pts)
        for f in faces:
            all_faces.append([idx + offset for idx in f])
        offset += len(pts)

    with open(filename, "w") as f:
        f.write("# vtk DataFile Version 3.0\n")
        f.write("Inner core best-fit planes\n")
        f.write("ASCII\n")
        f.write("DATASET POLYDATA\n")

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

        nfaces = len(all_faces)
        total_indices = nfaces * 5

        f.write(f"\nPOLYGONS {nfaces} {total_indices}\n")
        for face in all_faces:
            f.write("4 " + " ".join(map(str, face)) + "\n")

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

def main():

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

    df = load_data(sys.argv[1])

    best_lon = fit_two_plane(df)

    print("Best-fit planes:")
    print("Plane 1:", best_lon)
    print("Plane 2:", best_lon + 90)

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

    plane1 = build_plane_mesh(n1)
    plane2 = build_plane_mesh(n2)

    write_plane_vtk("frost_planes.vtk", [plane1, plane2])

    print("Wrote frost_planes.vtk")

if __name__ == "__main__":
    main()
