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

GRID_RES = 42     # 40–60 is ideal
IDW_POWER = 2.0
NEIGHBORS = 24

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

def sph_to_cart(lat, lon, r):
    lat = radians(lat)
    lon = radians(lon)
    return np.array([
        r*cos(lat)*cos(lon),
        r*cos(lat)*sin(lon),
        r*sin(lat)
    ])

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)

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

    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

# ------------------------------------------------------------
# Build inner-core point cloud
# ------------------------------------------------------------

def build_points(df):

    pts=[]
    vals=[]

    for la,lo,dep,dt in zip(df["TurnLa"],df["TurnLo"],df["TurnDep"],df["abs_dT"]):
        r=EARTH_RADIUS-dep
        p=sph_to_cart(la,lo,r)
        if np.linalg.norm(p)<=INNER_CORE_RADIUS+50:
            pts.append(p)
            vals.append(dt)

    return np.array(pts),np.array(vals)

# ------------------------------------------------------------
# IDW volumetric interpolation
# ------------------------------------------------------------

def interpolate_volume(points, values):

    grid=np.zeros((GRID_RES,GRID_RES,GRID_RES))

    xs=np.linspace(-INNER_CORE_RADIUS,INNER_CORE_RADIUS,GRID_RES)
    ys=xs
    zs=xs

    for i,x in enumerate(xs):
        print("Slice",i,"/",GRID_RES)
        for j,y in enumerate(ys):
            for k,z in enumerate(zs):

                p=np.array([x,y,z])
                r=np.linalg.norm(p)

                if r>INNER_CORE_RADIUS:
                    grid[i,j,k]=np.nan
                    continue

                dists=np.linalg.norm(points-p,axis=1)
                idx=np.argsort(dists)[:NEIGHBORS]

                d=dists[idx]
                v=values[idx]

                w=1/(d**IDW_POWER + 1e-6)
                grid[i,j,k]=np.sum(w*v)/np.sum(w)

    return grid

# ------------------------------------------------------------
# Write structured grid
# ------------------------------------------------------------

def write_structured_grid(filename,grid):

    nx,ny,nz=grid.shape

    with open(filename,"w") as f:
        f.write("# vtk DataFile Version 3.0\n")
        f.write("Inner core volumetric anisotropy\n")
        f.write("ASCII\n")
        f.write("DATASET STRUCTURED_POINTS\n")
        f.write(f"DIMENSIONS {nx} {ny} {nz}\n")
        f.write("ORIGIN -1220 -1220 -1220\n")
        f.write(f"SPACING {2440/(nx-1)} {2440/(ny-1)} {2440/(nz-1)}\n")

        f.write(f"\nPOINT_DATA {nx*ny*nz}\n")
        f.write("SCALARS anisotropy float 1\n")
        f.write("LOOKUP_TABLE default\n")

        for k in range(nz):
            for j in range(ny):
                for i in range(nx):
                    v=grid[i,j,k]
                    if np.isnan(v):
                        v=0
                    f.write(f"{v}\n")

# ------------------------------------------------------------
# Plane evolution as great circles
# ------------------------------------------------------------

def write_great_circle(filename,lon):

    n=plane_normal(lon)
    u=np.cross(n,[0,0,1])
    if np.linalg.norm(u)<1e-3:
        u=np.cross(n,[0,1,0])
    u=u/np.linalg.norm(u)
    v=np.cross(n,u)

    pts=[]
    for t in np.linspace(0,2*np.pi,200):
        p=INNER_CORE_RADIUS*(cos(t)*u+sin(t)*v)
        pts.append(p)

    with open(filename,"w") as f:
        f.write("# vtk DataFile Version 3.0\n")
        f.write("Great circle plane trace\n")
        f.write("ASCII\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"\nLINES 1 {len(pts)+1}\n")
        f.write(str(len(pts))+" "+" ".join(map(str,range(len(pts)))))

# ------------------------------------------------------------
# Fit plane per subset
# ------------------------------------------------------------

def fit_plane(df):

    pts=np.array([latlon_to_vec(la,lo) for la,lo in zip(df["TurnLa"],df["TurnLo"])])
    y=df["abs_dT"].values

    best_lon=0
    best_var=1e9

    for lon in range(0,180):
        n1=plane_normal(lon)
        n2=plane_normal(lon+90)

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

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

        if var<best_var:
            best_var=var
            best_lon=lon

    return best_lon

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

def main():

    df=load_data(sys.argv[1])
    points,values=build_points(df)

    print("Interpolating volume...")
    grid=interpolate_volume(points,values)

    print("Writing volume...")
    write_structured_grid("inner_core_volume.vtk",grid)

    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)
        write_great_circle(f"plane_trace_{i}.vtk",lon)
        print("subset",i,"plane:",lon)

if __name__=="__main__":
    main()
