#!/usr/bin/env python3

import numpy as np
import pandas as pd
import geopandas as gpd
from shapely.geometry import Point
from scipy.interpolate import griddata
import rasterio
from rasterio.transform import from_origin

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


def build_design_matrix(coords):
    x = coords[:,0]
    y = coords[:,1]
    z = coords[:,2]

    return np.column_stack([
        x, y, z,
        x*x, y*y, z*z,
        x*y, x*z, y*z
    ])


def fit_model(X, values):
    coeffs, _, _, _ = np.linalg.lstsq(X, values, rcond=None)
    predicted = X @ coeffs
    return coeffs, predicted


if __name__ == "__main__":

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

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

    signed_vals = df["dT"].values

    # 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

    valid = signed_vals != 0
    coords = coords[valid]
    signed_vals = signed_vals[valid]
    df = df.iloc[valid]

    # Fit degree-2 model
    X = build_design_matrix(coords)
    coeffs, predicted = fit_model(X, signed_vals)

    residual = signed_vals - predicted

    print("Exporting GeoPackage...")

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

    gdf = gpd.GeoDataFrame(
        {
            "observed_dT": signed_vals,
            "predicted_dT": predicted,
            "residual_dT": residual
        },
        geometry=geometry,
        crs="EPSG:4326"
    )

    gdf.to_file("degree2_model_points.gpkg", driver="GPKG")

    print("GeoPackage written: degree2_model_points.gpkg")

    # ---------------------------------------------------------
    # Optional: create raster of predicted field
    # ---------------------------------------------------------

    print("Exporting raster...")

    resolution = 1.0

    grid_lon = np.arange(-180, 181, resolution)
    grid_lat = np.arange(-90, 91, resolution)

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

    grid_pred = griddata(
        (df["TurnLo"].values, df["TurnLa"].values),
        predicted,
        (grid_lon_mesh, grid_lat_mesh),
        method="linear"
    )

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

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

    print("Raster written: degree2_model_predicted.tif")
