#!/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
OUTER_CORE_RADIUS = 3480.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 np.array([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)

# ------------------------------------------------------------
# Sphere mesh generator
# ------------------------------------------------------------

def build_sphere(radius, res=60):
    pts = []
    faces = []

    for i in range(res):
        theta = np.pi * i/(res-1)
        for j in range(res):
            phi = 2*np.pi*j/(res-1)

            x = radius*sin(theta)*cos(phi)
            y = radius*sin(theta)*sin(phi)
            z = radius*cos(theta)
            pts.append([x,y,z])

    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

# ------------------------------------------------------------
# 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 per subset
# ------------------------------------------------------------

def fit_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 range(0,180):
        n1 = plane_normal(lon)
        n2 = plane_normal(lon+90)

        x1 = np.abs(points @ n1)**2
        x2 = np.abs(points @ n2)**2

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

        if var < best_var:
            best_var = var
            best_lon = lon

    return best_lon

# ------------------------------------------------------------
# Write polydata
# ------------------------------------------------------------

def write_polydata(filename, pts, faces):

    with open(filename,"w") as f:
        f.write("# vtk DataFile Version 3.0\n")
        f.write("mesh\nASCII\nDATASET POLYDATA\n")

        f.write(f"POINTS {len(pts)} float\n")
        for p in pts:
            f.write(f"{p[0]} {p[1]} {p[2]}\n")

        f.write(f"\nPOLYGONS {len(faces)} {len(faces)*5}\n")
        for face in faces:
            f.write("4 " + " ".join(map(str,face))+"\n")

# ------------------------------------------------------------
# Write point cloud with vectors
# ------------------------------------------------------------

def write_points_with_vectors(filename, pts, vectors, scalars):

    n=len(pts)
    with open(filename,"w") as f:
        f.write("# vtk DataFile Version 3.0\n")
        f.write("points\nASCII\nDATASET POLYDATA\n")

        f.write(f"POINTS {n} float\n")
        for p in pts:
            f.write(f"{p[0]} {p[1]} {p[2]}\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")

        f.write("\nVECTORS ray_direction float\n")
        for v in vectors:
            f.write(f"{v[0]} {v[1]} {v[2]}\n")

        for name,data in scalars.items():
            f.write(f"\nSCALARS {name} float 1\nLOOKUP_TABLE default\n")
            for val in data:
                f.write(f"{val}\n")

# ------------------------------------------------------------
# Depth slicing frames
# ------------------------------------------------------------

def export_depth_timeseries(df):

    bins = np.linspace(df["TurnDep"].min(), df["TurnDep"].max(), 12)

    for i in range(len(bins)-1):

        sub = df[(df["TurnDep"]>=bins[i]) & (df["TurnDep"]<bins[i+1])]

        pts=[]
        vecs=[]

        for la,lo,dep in zip(sub["TurnLa"],sub["TurnLo"],sub["TurnDep"]):
            r = EARTH_RADIUS - dep
            p = sph_to_cart(la,lo,r)
            pts.append(p)

            v = latlon_to_vec(la,lo)
            vecs.append(v)

        write_points_with_vectors(
            f"depth_frame_{i:02d}.vtk",
            pts,
            vecs,
            {"abs_dT":sub["abs_dT"].values}
        )

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

def main():

    df=load_data(sys.argv[1])

    # 1) inner core sphere
    pts,faces=build_sphere(INNER_CORE_RADIUS)
    write_polydata("inner_core_surface.vtk",pts,faces)

    # 2) outer core boundary
    pts,faces=build_sphere(OUTER_CORE_RADIUS)
    write_polydata("outer_core_boundary.vtk",pts,faces)

    # 3) turning points + ray vectors
    pts=[]
    vecs=[]

    for la,lo,dep in zip(df["TurnLa"],df["TurnLo"],df["TurnDep"]):
        r=EARTH_RADIUS-dep
        pts.append(sph_to_cart(la,lo,r))
        vecs.append(latlon_to_vec(la,lo))

    write_points_with_vectors(
        "turning_points_vectors.vtk",
        pts,
        vecs,
        {"abs_dT":df["abs_dT"].values,
         "depth":df["TurnDep"].values}
    )

    # 4) depth animation
    export_depth_timeseries(df)

    # 5) subset plane time stack
    subsets=[
        df,
        df[df["TurnDep"]>6000],
        df[df["Ksi"]<=35],
        df[(df["TurnDep"]>6000)&(df["Ksi"]<=35)]
    ]

    for i,sub in enumerate(subsets):
        lon=fit_plane(sub)
        print(f"subset {i} best plane:",lon)

if __name__=="__main__":
    main()
