"""FV3GFS preprocessing interfaces used by the ACE model package."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np


@dataclass(frozen=True)
class VerticalCoordinate:
    a_pa: np.ndarray
    b: np.ndarray
    source_indices: np.ndarray

    @classmethod
    def ace_default(cls) -> "VerticalCoordinate":
        return cls(
            np.asarray([64.247, 5167.14603, 12905.42546, 13982.4677, 12165.28766, 8910.07678, 4955.72632, 2155.78385, 0.0]),
            np.asarray([0.0, 0.0, 0.01755, 0.11746, 0.2896, 0.49806, 0.72625, 0.88192, 1.0]),
            np.asarray([0, 18, 26, 31, 36, 41, 47, 53, 63]),
        )

    def interface_pressure(self, surface_pressure: np.ndarray) -> np.ndarray:
        ps = np.asarray(surface_pressure)
        if ps.ndim < 2:
            raise ValueError("surface_pressure must end in [lat, lon]")
        return np.einsum("k,...ij->...kij", self.b, ps) + self.a_pa.reshape((1,) * (ps.ndim - 2) + (9, 1, 1))


def conservative_vertical_coarsen(values: np.ndarray, layer_thickness: np.ndarray, coordinate: VerticalCoordinate | None = None, layer_axis: int = -3) -> tuple[np.ndarray, np.ndarray]:
    coord = coordinate or VerticalCoordinate.ace_default()
    x = np.moveaxis(np.asarray(values, dtype=np.float64), layer_axis, -3)
    dp = np.moveaxis(np.asarray(layer_thickness, dtype=np.float64), layer_axis, -3)
    if x.shape[-3] != 63 or dp.shape != x.shape:
        raise ValueError(f"expected values/dp shape [...,63,H,W], got {x.shape} and {dp.shape}")
    means, thickness = [], []
    for start, end in zip(coord.source_indices[:-1], coord.source_indices[1:]):
        chunk_dp, chunk_x = dp[..., start:end, :, :], x[..., start:end, :, :]
        dpk = chunk_dp.sum(axis=-3)
        means.append((chunk_x * chunk_dp).sum(axis=-3) / np.maximum(dpk, 1e-12))
        thickness.append(dpk)
    return np.moveaxis(np.stack(means, axis=-3), -3, layer_axis), np.moveaxis(np.stack(thickness, axis=-3), -3, layer_axis)


def conservative_regrid(fields: np.ndarray, regrid_fn: Callable[[np.ndarray], np.ndarray] | None = None) -> np.ndarray:
    if regrid_fn is None:
        raise NotImplementedError("MISSING: NOAA fregrid or another conservative cubed-sphere adapter")
    return np.asarray(regrid_fn(np.asarray(fields)))


def spherical_harmonic_roundtrip(fields: np.ndarray, filter_fn: Callable[[np.ndarray], np.ndarray] | None = None) -> np.ndarray:
    return np.asarray(fields if filter_fn is None else filter_fn(np.asarray(fields))).copy()