"""
The unstructured physics grid.

LMDZ physics variables are stored on a flattened `physical_points` dimension.
The values sit exactly on a regular latitude/longitude mesh, so rebuilding the
2D grid is a scatter, not an interpolation.

Two things moved out of here during the refactor: radian-to-degree conversion is
now done once, for the whole dataset, by conventions.py, and the scatter indices
are computed once per file rather than once per plot.
"""

from dataclasses import dataclass
from functools import lru_cache

import numpy as np


@dataclass(frozen=True)
class UnstructuredGrid:
    """
    The mapping from flattened physics points onto a regular mesh.
    """
    dim: str
    lats: np.ndarray          # unique latitudes, ascending
    lons: np.ndarray          # unique longitudes, ascending, seam column appended
    row: np.ndarray           # target row of each physics point
    col: np.ndarray           # target column of each physics point
    seam_source: int = None   # column duplicated to close the -180/+180 seam
    exact: bool = True        # False when the tolerant lookup was needed

    @property
    def shape(self):
        return self.lats.size, self.lons.size


def build_grid(ds, info):
    """
    Build the scatter mapping for one dataset.

    `info` is the dict conventions.normalize() recorded as `report.unstructured`,
    naming the dimension and its latitude/longitude variables.
    """
    if not info:
        return None
    lat_values = np.asarray(ds[info['lat']].values, dtype=float).ravel()
    lon_values = np.asarray(ds[info['lon']].values, dtype=float).ravel()

    uniq_lats = np.unique(lat_values)
    uniq_lons = np.unique(lon_values)

    row = np.searchsorted(uniq_lats, lat_values)
    col = np.searchsorted(uniq_lons, lon_values)

    exact = (
        not np.any(row >= uniq_lats.size) and not np.any(col >= uniq_lons.size)
        and np.array_equal(uniq_lats[np.clip(row, 0, uniq_lats.size - 1)], lat_values)
        and np.array_equal(uniq_lons[np.clip(col, 0, uniq_lons.size - 1)], lon_values)
    )
    if not exact:
        # Should not happen, but fall back on a tolerant nearest-point lookup
        print("Warning: physical_points coordinates did not match exactly; using tolerant lookup.")
        row = np.array([np.abs(uniq_lats - v).argmin() for v in lat_values])
        col = np.array([np.abs(uniq_lons - v).argmin() for v in lon_values])

    seam_source = None
    if np.any(np.isclose(uniq_lons, -180.0)):
        seam_source = int(np.where(np.isclose(uniq_lons, -180.0))[0][0])
        uniq_lons = np.append(uniq_lons, 180.0)

    return UnstructuredGrid(dim=info['dim'], lats=uniq_lats, lons=uniq_lons,
                            row=row, col=col, seam_source=seam_source, exact=exact)


def to_grid(data, grid):
    """
    Scatter values from the flattened grid onto (..., nlat, nlon).

    The physics dimension must be the last axis of `data`; any leading axes are
    carried through untouched.
    """
    data = np.asarray(data, dtype=float)
    nlat = grid.lats.size
    nlon = grid.lons.size - (1 if grid.seam_source is not None else 0)

    out = np.full(data.shape[:-1] + (nlat, nlon), np.nan)
    out[..., grid.row, grid.col] = data

    # A pole is stored as a single physics point; spread it across all longitudes
    # so the map does not show one coloured cell and a row of holes.
    for index in (0, -1):
        row = out[..., index, :]
        good = np.isfinite(row)
        single = good.sum(axis=-1) == 1
        if np.any(single):
            values = np.nansum(np.where(good, row, 0.0), axis=-1)
            row[single, :] = values[single, None] if row.ndim > 1 else values
            out[..., index, :] = row

    if grid.seam_source is not None:
        # Duplicate the -180 column at +180 so the map closes
        out = np.concatenate([out, out[..., [grid.seam_source]]], axis=-1)

    return out