#!/usr/bin/env python3

import numpy as np
import pandas as pd

CACHE_DIR = "analysis_cache"
DATAFILE = "frost2021.dat"
CLUSTER_PERCENTILE = 85


def vec_to_latlon(v):
    v = v / np.linalg.norm(v)
    lat = np.degrees(np.arcsin(v[2]))
    lon = np.degrees(np.arctan2(v[1], v[0]))
    return lat, lon


def compute_tensor(coords):
    M = np.zeros((3,3))
    for x in coords:
        M += np.outer(x, x)
    return M


def eigen_analysis(M):
    eigvals, eigvecs = np.linalg.eigh(M)
    idx = np.argsort(eigvals)
    eigvals = eigvals[idx]
    eigvecs = eigvecs[:, idx]
    return eigvals, eigvecs


if __name__ == "__main__":

    with open(DATAFILE) 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(
        DATAFILE,
        sep=r"\s+",
        engine="python",
        skiprows=header_idx
    )

    print("Available columns:")
    print(df.columns.tolist())

    # Resolve dT
    df["dT"] = df["dTabdf"].copy()
    mask = df["dT"] == 0
    df.loc[mask, "dT"] = df.loc[mask, "dTbcdf"]
    df["abs_dT"] = df["dT"].abs()

    # ---- Identify correct depth column ----
    depth_column = None
    for col in df.columns:
        if "Dep" in col:
            depth_column = col
            break

    if depth_column is None:
        raise ValueError("No depth column found.")

    print("Using depth column:", depth_column)

    # Convert to vectors
    lat = np.radians(df["TurnLa"].values)
    lon = np.radians(df["TurnLo"].values)

    coords = np.vstack((
        np.cos(lat) * np.cos(lon),
        np.cos(lat) * np.sin(lon),
        np.sin(lat)
    )).T

    values = df["abs_dT"].values
    depths = df[depth_column].values

    # Cluster mask
    threshold = np.percentile(values, CLUSTER_PERCENTILE)
    cluster_mask = values >= threshold

    coords_cluster = coords[cluster_mask]
    depths_cluster = depths[cluster_mask]

    print("\nTotal cluster points:", len(coords_cluster))

    # Sort by depth
    idx_sort = np.argsort(depths_cluster)
    coords_sorted = coords_cluster[idx_sort]
    depths_sorted = depths_cluster[idx_sort]

    N = len(coords_sorted)
    bin_size = N // 3

    bins = {
        "Shallow": coords_sorted[:bin_size],
        "Mid": coords_sorted[bin_size:2*bin_size],
        "Deep": coords_sorted[2*bin_size:]
    }

    for name, coords_bin in bins.items():

        M = compute_tensor(coords_bin)
        eigvals, eigvecs = eigen_analysis(M)

        lat_axis, lon_axis = vec_to_latlon(eigvecs[:, -1])
        spread = eigvals[-1] - eigvals[0]

        print(f"\n{name} bin:")
        print(" Points:", len(coords_bin))
        print(" Axis lat:", lat_axis, "lon:", lon_axis)
        print(" Spread:", spread)
