#!/usr/bin/env python3

import sys
import numpy as np
import pandas as pd
from math import radians, sin, cos, asin
from scipy.interpolate import griddata
import rasterio
from rasterio.transform import from_origin
import geopandas as gpd
from shapely.geometry import Point

# ------------------------------------------------------------
# Geometry helpers (identical logic to frost_visualizations.py)
# ------------------------------------------------------------

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

# ------------------------------------------------------------
# Compute residuals
# ------------------------------------------------------------

def compute_residuals(df, best_lon, coeff):

    points = np.array([
        latlon_to_vec(la, lo)
        for la, lo in zip(df["TurnLa"], df["TurnLo"])
    ])

    n1 = plane_normal(best_lon)
    n2 = plane_normal(best_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

    pred = coeff[0] + coeff[1]*x1 + coeff[2]*x2
    residuals = df["abs_dT"].values - pred

    return residuals

# ------------------------------------------------------------
# Export GeoTIFF overlay
# ------------------------------------------------------------

def export_geotiff(df, residuals, outfile="residual_overlay.tif", resolution=1.0):

    lons = df["TurnLo"].values
    lats = df["TurnLa"].values

    # Create regular global grid
    grid_lon = np.arange(-180, 180 + resolution, resolution)
    grid_lat = np.arange(-90, 90 + resolution, resolution)

    grid_lon_mesh, grid_lat_mesh = np.meshgrid(grid_lon, grid_lat)

    # Interpolate residuals to grid
    grid_res = griddata(
        (lons, lats),
        residuals,
        (grid_lon_mesh, grid_lat_mesh),
        method="linear"
    )

    # Define geotransform
    transform = from_origin(
        west=-180,
        north=90,
        xsize=resolution,
        ysize=resolution
    )

    with rasterio.open(
        outfile,
        "w",
        driver="GTiff",
        height=grid_res.shape[0],
        width=grid_res.shape[1],
        count=1,
        dtype=grid_res.dtype,
        crs="EPSG:4326",
        transform=transform,
    ) as dst:
        dst.write(np.flipud(grid_res), 1)

    print(f"GeoTIFF written to {outfile}")

# ------------------------------------------------------------
# Export vector point layer (GeoPackage)
# ------------------------------------------------------------

def export_geopackage(df, residuals, outfile="residual_points.gpkg"):

    geometry = [
        Point(lon, lat)
        for lon, lat in zip(df["TurnLo"], df["TurnLa"])
    ]

    gdf = gpd.GeoDataFrame(
        df.assign(residual=residuals),
        geometry=geometry,
        crs="EPSG:4326"
    )

    gdf.to_file(outfile, driver="GPKG")

    print(f"GeoPackage written to {outfile}")

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

def main():

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

    df = load_data(sys.argv[1])
    best_lon, coeff = fit_two_plane(df)

    print("Best plane:", best_lon)
    print("Orthogonal plane:", (best_lon+90)%180)

    residuals = compute_residuals(df, best_lon, coeff)

    export_geotiff(df, residuals)
    export_geopackage(df, residuals)


if __name__ == "__main__":
    main()
