from __future__ import annotations
import numpy as np
from .transform import var_transform


def build_feature_cube(ds, clim, feature_vars, tidx, months, lat2d, lon2d, reg, include_lat_lon=True, include_region_feature=False, precip_floor=1.0):
    cubes = []
    names = []
    for v in feature_vars:
        cubes.append(var_transform(ds, clim, v, tidx, months, precip_floor=precip_floor))
        names.append(v + ('_factor' if v == 'pr' else '_ano'))
    ntime = len(tidx)
    if include_lat_lon:
        lat = np.broadcast_to(lat2d[None, :, :], (ntime, lat2d.shape[0], lat2d.shape[1])).astype(np.float32)
        lon = np.broadcast_to(lon2d[None, :, :], (ntime, lon2d.shape[0], lon2d.shape[1])).astype(np.float32)
        cubes.extend([lat, lon])
        names.extend(['lat', 'lon'])
    if include_region_feature:
        rr = np.broadcast_to(reg[None, :, :], (ntime, reg.shape[0], reg.shape[1])).astype(np.float32)
        cubes.append(rr)
        names.append('reg')
    X = np.stack(cubes, axis=-1).astype(np.float32)
    return X, names


def select_samples(Xcube: np.ndarray, ycube: np.ndarray, smask: np.ndarray):
    X = Xcube[:, smask, :].reshape(-1, Xcube.shape[-1])
    y = ycube[:, smask].reshape(-1)
    ok = np.isfinite(y) & np.all(np.isfinite(X), axis=1)
    return X[ok], y[ok]


def select_samples_with_aux(Xcube: np.ndarray, ycube: np.ndarray, smask: np.ndarray, auxcube: np.ndarray | None = None):
    X = Xcube[:, smask, :].reshape(-1, Xcube.shape[-1])
    y = ycube[:, smask].reshape(-1)
    ok = np.isfinite(y) & np.all(np.isfinite(X), axis=1)
    aux = None
    if auxcube is not None:
        aux = auxcube[:, smask].reshape(-1)
        ok = ok & np.isfinite(aux)
    if aux is None:
        return X[ok], y[ok], None
    return X[ok], y[ok], aux[ok]


def predict_to_grid(model, Xcube: np.ndarray, smask: np.ndarray) -> np.ndarray:
    out = np.full(Xcube.shape[:3], np.nan, dtype=np.float32)
    X = Xcube[:, smask, :].reshape(-1, Xcube.shape[-1])
    ok = np.all(np.isfinite(X), axis=1)
    flat = np.full(X.shape[0], np.nan, dtype=np.float32)
    if np.any(ok):
        flat[ok] = model.predict(X[ok]).astype(np.float32)
    out[:, smask] = flat.reshape(Xcube.shape[0], int(np.sum(smask)))
    return out


def spatial_masks(valid_mask: np.ndarray, reg: np.ndarray, spatial_mode: str, valid_regions=(1,2,3,4,5,6,7)):
    if spatial_mode == 'national':
        yield 'national', valid_mask
    elif spatial_mode == 'regional':
        for r in valid_regions:
            yield f'R{int(r)}', valid_mask & (reg == int(r))
    else:
        raise ValueError(f'Unknown spatial_mode={spatial_mode}')