import numpy as np

from util.setup import home_dir
from scipy.interpolate import interp2d
from scipy import linalg
from scipy.spatial import KDTree
from sklearn.neighbors import BallTree
from util.util import haversine_np
from scipy.ndimage import gaussian_filter


class MyGenericException(Exception):
    def __init__(self, message):
        self.message = message


# Author: T.Rink
# For nearest neighbor/interpolation of a function whose domain is a geostationary satellite fixed grid.
# This is a utility wrapper for sklearn.neighbors.BallTree adapted for 2D manifold (Longitude, Latitude).
class LonLatGrid:
    # grd_lons, grd_lats: the longitude, latitude of each grid point (must be 2D grids), must have same shape
    # Incoming longitude must be in range: 0 - 360 degrees
    # Can have NaN for off Earth grid points (these are handled internally).
    # closeness_threshold (m): if < distance of located_point to target, return off grid
    # reduce: resolution blow down factor
    # leaf_size: Number of points at which to switch to brute-force, default=40
    def __init__(self, grd_lons, grd_lats, closeness_threshold=2000, reduce=1, leaf_size=40):
        if grd_lons.shape != grd_lats.shape:
            raise MyGenericException('incoming lons,lats must have same shape')

        self.grd_lons = grd_lons
        self.grd_lats = grd_lats
        self.grd_lons = grd_lons[::reduce, ::reduce]
        self.grd_lats = grd_lats[::reduce, ::reduce]
        self.shape = self.grd_lons.shape
        self.lenx = self.shape[1]
        self.leny = self.shape[0]
        self.fnc_dct = {}
        self.x = np.zeros(9, np.float32)
        self.y = np.zeros(9, np.float32)
        self.z = np.zeros(9, np.float32)

        flons = self.grd_lons.flatten()
        flats = self.grd_lats.flatten()

        keep = np.invert(np.isnan(flons))

        self.map_indexes = (np.arange(self.grd_lons.size))[keep]

        flons = flons[keep]
        flats = flats[keep]
        self.lon_lo, self.lon_hi = flons.min(), flons.max()
        self.lat_lo, self.lat_hi = flats.min(), flats.max()

        points = np.stack([flons, flats], axis=1)

        self.kd = BallTree(np.deg2rad(points), leaf_size=leaf_size, metric='haversine')
        self.closeness_threshold = closeness_threshold * reduce
        self.reduce = reduce

    # locate nearest neighbor for incoming target in earth coordinates (Should not have NaNs)
    # lons, lats can be flat or 2D
    # incoming longitude must be in range: 0 - 360 degrees
    # returns NN indexes relative to the grid determined in the ctr.
    def value_to_index(self, lons, lats):
        if lons.shape != lats.shape:
            raise MyGenericException('incoming lons,lats must have same shape')

        if len(lons) == 0 or len(lats) == 0:
            return None, None

        lons = lons.flatten()
        lats = lats.flatten()

        xy = np.stack([lons, lats], axis=1)

        # Assumes nn for a query pont returned in order of increasing distance
        dist, indices = self.kd.query(np.deg2rad(xy), k=1)
        dist = dist[:, 0]
        indices = indices[:, 0]
        dist *= 6370000  # convert unit radius to meters

        valid = (indices < self.map_indexes.size) & (dist < self.closeness_threshold)
        indices = indices[valid]

        return self.map_indexes[indices], indices

    def earth_to_lc_s(self, lons, lats):
        if lons.shape != lats.shape:
            raise MyGenericException('incoming lons,lats must have same shape')

        if len(lons) == 0 or len(lats) == 0:
            return None, None

        lons = np.where(lons < 0, lons + 360, lons)

        lons = lons.flatten()
        lats = lats.flatten()

        xy = np.stack([lons, lats], axis=1)

        # Assumes nn for a query pont returned in order of increasing distance
        dist, indices = self.kd.query(np.deg2rad(xy), k=1)
        dist = dist[:, 0]
        indices = indices[:, 0]
        dist *= 6370000  # convert unit radius to meters

        valid = (indices < self.map_indexes.size) & (dist < self.closeness_threshold)
        indices = np.where(valid, indices, -1)
        ll = np.where(indices >= 0, (indices / self.lenx).astype(np.int32), indices)
        cc = np.where(indices >= 0, indices % self.lenx, indices)

        cc *= self.reduce
        ll *= self.reduce

        return cc, ll

    # bi-linear interpolation of a function on the grid, the range of which is grd_zvals for target lons,lats.
    # returns the interpolated values at each target point.
    # target lons, lats should not have NaNs, and lons must be in range 0 - 360 degrees
    def interp(self, grid_zvals, lons, lats):
        if grid_zvals.shape != self.shape:
            raise MyGenericException('range and domain must have same 2D shape')

        if lons.shape != lats.shape:
            raise MyGenericException('incoming lons, lats must have same shape')

        abs_idxs, rel_idxs = self.value_to_index(lons, lats)

        x = np.zeros(9, np.float32)
        y = np.zeros(9, np.float32)
        z = np.zeros(9, np.float32)

        lons = lons.flatten()
        lats = lats.flatten()

        outv = np.zeros(lons.size, np.float32)
        outv.fill(np.nan)

        for k in range(len(abs_idxs)):
            idx = abs_idxs[k]
            lon = lons[k]
            lat = lats[k]

            fnc = self.fnc_dct.get(idx)
            if fnc is None:
                gj = int(idx / self.lenx)
                gi = int(idx % self.lenx)

                L = gi-1
                R = gi+1
                D = gj-1
                U = gj+1

                x[0] = self.grd_lons[D, L]
                y[0] = self.grd_lats[D, L]
                z[0] = grid_zvals[D, L]
                x[1] = self.grd_lons[D, gi]
                y[1] = self.grd_lats[D, gi]
                z[1] = grid_zvals[D, gi]
                x[2] = self.grd_lons[D, R]
                y[2] = self.grd_lats[D, R]
                z[2] = grid_zvals[D, R]
                x[3] = self.grd_lons[gj, L]
                y[3] = self.grd_lats[gj, L]
                z[3] = grid_zvals[gj, L]
                x[4] = self.grd_lons[gj, gi]
                y[4] = self.grd_lats[gj, gi]
                z[4] = grid_zvals[gj, gi]
                x[5] = self.grd_lons[gj, R]
                y[5] = self.grd_lats[gj, R]
                z[5] = grid_zvals[gj, R]
                x[6] = self.grd_lons[U, L]
                y[6] = self.grd_lats[U, L]
                z[6] = grid_zvals[U, L]
                x[7] = self.grd_lons[U, gi]
                y[7] = self.grd_lats[U, gi]
                z[7] = grid_zvals[U, gi]
                x[8] = self.grd_lons[U, R]
                y[8] = self.grd_lats[U, R]
                z[8] = grid_zvals[U, R]

                if np.isnan(x).sum() > 0:
                    continue
                if np.isnan(y).sum() > 0:
                    continue

                fnc = interp2d(x, y, z, copy=False)
                self.fnc_dct[idx] = fnc

            outv[k] = fnc(lon, lat)

        return outv

    def check_inside(self, lon, lat):
        if lon < 0:
            lon += 360
        return (self.lon_lo < lon < self.lon_hi) & (self.lat_lo < lat < self.lat_hi)


DEG_TO_RAD = np.pi/180.0;
RAD_TO_DEG = 180.0/np.pi;


r_pol = 6356.5838  # km
r_eq = 6378.1690  # km
# AHI -------------------------------
# h = 42164.0  # barycentric height, km
# sub_lon = 140.7
# sub_lon *= DEG_TO_RAD
# scan_geom = 'GEOS'
# signs modified from static nav file to work with ABI FGF code
# CFAC = 5.588799E-05
# LFAC = -5.588799E-05
# COFF = -0.1537199
# LOFF = 0.1537199

# computed from static nav file (lon, lat) grid
# Note 0-based to the pixel center
# CFAC = 5.58924125e-05
# LFAC = -5.58810490e-05
# COFF = -1.53678977e-01
# LOFF = 1.53644345e-01

# GOES ------------------------------
h = 42164.16 # barycentric height, km
scan_geom = 'GOES'
sub_lon = -75.0
sub_lon *= DEG_TO_RAD

# official for FD
CFAC = 5.6E-05
LFAC = -5.6E-05
COFF = -0.151844
LOFF = 0.151844

# official for CONUS
#CFAC = 5.6E-05
#COFF = -0.101332
#LFAC = -5.6E-05
#LOFF = 0.128212

# computed for CLAVRx FD
#CFAC = 5.60016368e-05
#LFAC = -5.59941969e-05
#COFF = -1.51780260e-01
#LOFF = 1.51773560e-01

# 65536 = 2^16


def earth_to_sat(geographic_lon, geographic_lat):

    geographic_lon *= DEG_TO_RAD
    geographic_lat *= DEG_TO_RAD

    geocentric_lat = np.arctan(((r_pol*r_pol)/(r_eq*r_eq))*np.tan(geographic_lat))

    r_earth = r_pol / np.sqrt(1.0 - ((r_eq * r_eq - r_pol * r_pol) / (r_eq * r_eq)) * np.cos(geocentric_lat) * np.cos(geocentric_lat))

    r_1 = h - r_earth * np.cos(geocentric_lat) * np.cos(geographic_lon - sub_lon)
    r_2 = -r_earth * np.cos(geocentric_lat) * np.sin(geographic_lon - sub_lon)
    r_3 = r_earth * np.sin(geocentric_lat)

    if r_1 > h:
        return np.nan, np.nan

    if scan_geom == 'GEOS':
        lamda_sat = np.arctan(-r_2/r_1)
        theta_sat = np.arcsin(r_3/np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3))
    elif scan_geom == 'GOES':
        lamda_sat = np.arcsin(-r_2/np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3))
        theta_sat = np.arctan(r_3/r_1)

    return lamda_sat, theta_sat


def earth_to_sat_s(geographic_lon, geographic_lat):

    geographic_lon = geographic_lon * DEG_TO_RAD
    geographic_lat = geographic_lat * DEG_TO_RAD

    geocentric_lat = np.arctan(((r_pol*r_pol)/(r_eq*r_eq))*np.tan(geographic_lat))

    r_earth = r_pol / np.sqrt(1.0 - ((r_eq * r_eq - r_pol * r_pol) / (r_eq * r_eq)) * (np.cos(geocentric_lat))**2)

    r_1 = h - r_earth * np.cos(geocentric_lat) * np.cos(geographic_lon - sub_lon)
    r_2 = -r_earth * np.cos(geocentric_lat) * np.sin(geographic_lon - sub_lon)
    r_3 = r_earth * np.sin(geocentric_lat)

    if scan_geom == 'GEOS':
        lamda_sat = np.arctan(-r_2/r_1)
        theta_sat = np.arcsin(r_3/np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3))
    elif scan_geom == 'GOES':
        lamda_sat = np.arcsin(-r_2/np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3))
        theta_sat = np.arctan(r_3/r_1)

    np.where(r_1 > h, lamda_sat == np.nan, lamda_sat)
    np.where(r_1 > h, theta_sat == np.nan, theta_sat)

    return lamda_sat, theta_sat


# def earth_to_sat_s(geographic_lons, geographic_lats):
#     lambdas = []
#     thetas = []
#
#     for i in range(len(geographic_lons)):
#         l, t = earth_to_sat(geographic_lons[i], geographic_lats[i])
#         lambdas.append(l)
#         thetas.append(t)
#
#     return lambdas, thetas


def sat_to_lc(lamda, theta):
    # c = COFF + np.rint(lamda * (1/np.power(2, 16)) * CFAC)
    # l = LOFF + np.rint(theta * (1/np.power(2, 16)) * LFAC)

    # (float(2 ** 16) * (float(lc) - off)) / float(fac) * DEG_TO_RAD

    cc = (lamda - COFF) / CFAC
    ll = (theta - LOFF) / LFAC

    cc = np.floor(cc + 0.5)
    ll = np.floor(ll + 0.5)

    # cc = int(cc)
    # ll = int(ll)

    cc = cc.astype(np.int32)
    ll = ll.astype(np.int32)

    return cc, ll


def test(flons, flats):
    fflons = flons.flatten()
    fflats = flats.flatten()
    num = len(fflons)

    for i in range(0, num-1, 50000):
        lon = fflons[i]
        lat = fflats[i]
        if lon == -999.0:
            continue
        c, l = earth_to_lc(lon, lat)
        print(lon, flons[l,c])
        print(lat, flats[l,c])
        print(haversine_np(lon, lat, flons[l,c], flats[l,c]))
        print('-----------------------')


def earth_to_lc(lon, lat):
    lamda, theta = earth_to_sat(lon, lat)
    if np.isnan(lamda):
        return None, None
    cc, ll = sat_to_lc(lamda, theta)
    return cc, ll


# def earth_to_lc_s(lons, lats):
#     num = lons.shape[0]
#     cc_s = []
#     ll_s = []
#     for i in range(num):
#         cc, ll = earth_to_lc(lons[i], lats[i])
#         if cc is None:
#             cc_s.append(-1)
#             ll_s.append(-1)
#         else:
#             cc_s.append(cc)
#             ll_s.append(ll)
#
#     return cc_s, ll_s


def earth_to_lc_s(lons, lats):
    lamda, theta = earth_to_sat_s(lons, lats)
    cc, ll = sat_to_lc(lamda, theta)
    np.where(np.isnan(cc), cc == -1, cc)
    np.where(np.isnan(ll), ll == -1, ll)
    return cc, ll


def earth_to_indexs(lons, lats, len_x):
    num = lons.shape[0]
    idxs = []
    for i in range(num):
        cc, ll = earth_to_lc(lons[i], lats[i])
        if cc is None:
            idxs.append(-1)
        else:
            idx = ll * len_x + cc
            idxs.append(idx)
    idxs = np.array(idxs)

    return idxs


def compute_scale_offset(lon_a, lat_a, col_a, line_a, lon_b, lat_b, col_b, line_b):
    lamda_a, theta_a = earth_to_sat(lon_a, lat_a)
    lamda_b, theta_b = earth_to_sat(lon_b, lat_b)

    # To setup the problem below
    # (CFAC * col_a) + COFF = lamda_a
    # (CFAC * col_b) + COFF = lamda_b

    # (LFAC * line_a) + LOFF = theta_a
    # (LFAC * line_b) + LOFF = theta_b

    a = np.array([[col_a, 1], [col_b, 1]])
    b = np.array([lamda_a, lamda_b])
    x = linalg.solve(a, b)
    print(x)

    a = np.array([[line_a, 1], [line_b, 1]])
    b = np.array([theta_a, theta_b])
    x = linalg.solve(a, b)
    print(x)