import numpy as np
from numpy import ndarray
from sklearn.neighbors import KNeighborsRegressor

__all__ = ["GridSampler"]


class GridSampler(object):
    def __init__(self, data: ndarray, longitude: ndarray, latitude: ndarray, k: int = 5):
        """
        :param data:            dbz
        :param longitude:       经度
        :param latitude:        维度
        :param k:               取几个临近点
        ※※※※※※※※※※※※※※※※※※※※
        ※※                             ※※
        ※※    所有的参数都是一维的数据   ※※
        ※※                             ※※
        ※※※※※※※※※※※※※※※※※※※※
        """
        assert data.ndim == 1 and longitude.ndim == 1 and latitude.ndim == 1 and k > 1

        longitude, self.lon_avg, self.lon_std = _normalize(longitude)
        latitude, self.lat_avg, self.lat_std = _normalize(latitude)

        inputs = np.dstack([longitude, latitude])[0]
        self.knn = KNeighborsRegressor(n_neighbors=k, weights="distance")
        data[np.isnan(data)] = 0  # 雷达缺失值用 0 填充，否则 KNN 会报错
        self.knn.fit(inputs, data)

    def map_data(self, lon: ndarray, lat: ndarray) -> ndarray:
        shape = lon.shape
        lon = lon.ravel()
        lat = lat.ravel()

        lon = (lon - self.lon_avg) / self.lon_std
        lat = (lat - self.lat_avg) / self.lat_std

        inputs = np.dstack([lon, lat])[0]
        outputs = self.knn.predict(inputs).reshape(shape)

        return outputs


def _normalize(data: ndarray):
    avg = data.mean()
    std = data.std()
    data = (data - avg) / std
    return data, avg, std
