import numpy as np

from scipy import linalg
from util.util import haversine_np

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
invf = 298.2572221  # inverse flattening
f = 1.0/invf
fp = 1.0/((1.0-f)*(1.0-f))
d = h*h - r_eq*r_eq
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 goes_to_geos(lamda_goes, theta_goes):
    theta_geos = np.arcsin(np.sin(theta_goes) * np.cos(lamda_goes))
    lamda_geos = np.arctan(np.tan(lamda_goes) / np.cos(theta_goes))

    return lamda_geos, theta_geos


class GEOSNavigation:
    def __init__(self, sub_lon=-75.0, barycentric_height=42164.16, scan_geom='GOES', CFAC=5.6E-05, LFAC=-5.6E-05, COFF=-0.151844, LOFF=0.151844, num_elems=5424, num_lines=5424):
        self.sub_lon = sub_lon * DEG_TO_RAD
        self.h = barycentric_height
        self.scan_geom = scan_geom
        self.CFAC = CFAC
        self.LFAC = LFAC
        self.COFF = COFF
        self.LOFF = LOFF
        self.num_lines = num_lines
        self.num_elems = num_elems

    def earth_to_sat(self, 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 = self.h - r_earth * np.cos(geocentric_lat) * np.cos(geographic_lon - self.sub_lon)
        r_2 = -r_earth * np.cos(geocentric_lat) * np.sin(geographic_lon - self.sub_lon)
        r_3 = r_earth * np.sin(geocentric_lat)
        r = np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3)

        if r >= self.h:
            return np.nan, np.nan

        if self.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 self.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(self, 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 = self.h - r_earth * np.cos(geocentric_lat) * np.cos(geographic_lon - self.sub_lon)
        r_2 = -r_earth * np.cos(geocentric_lat) * np.sin(geographic_lon - self.sub_lon)
        r_3 = r_earth * np.sin(geocentric_lat)
        r = np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3)

        if self.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 self.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)

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

        return lamda_sat, theta_sat

    def sat_to_lc(self, 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 - self.COFF) / self.CFAC
        ll = (theta - self.LOFF) / self.LFAC

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

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

        return cc, ll

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

    def earth_to_lc_s(self, lons, lats):
        lamda, theta = self.earth_to_sat_s(lons, lats)
        cc, ll = self.sat_to_lc(lamda, theta)

        cc = np.where(np.isnan(lamda), -1, cc)
        cc = np.where(cc < 0, -1, cc)
        cc = np.where(cc >= self.num_elems, -1, cc)

        ll = np.where(np.isnan(theta), -1, ll)
        ll = np.where(ll < 0, -1, ll)
        ll = np.where(ll >= self.num_lines, -1, ll)

        cc[ll == -1] = -1
        ll[cc == -1] = -1

        return cc, ll

    def lc_to_earth(self, cc, ll):
        x = cc * self.CFAC + self.COFF
        y = ll * self.LFAC + self.LOFF

        lon, lat = self.sat_to_earth_s(x, y)

        return lon, lat

    def earth_to_indexs(self, lons, lats, len_x):
        num = lons.shape[0]
        idxs = []
        for i in range(num):
            cc, ll = self.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 sat_to_earth_s(self, x, y):
        if self.scan_geom == 'GOES':
            x, y = goes_to_geos(x, y)

        c1 = (h * np.cos(x) * np.cos(y)) * (h * np.cos(x) * np.cos(y))
        c2 = (np.cos(y) * np.cos(y) + fp * np.sin(y) * np.sin(y)) * d

        c1 = np.where(c1 < c2, np.nan, c1)
        c2 = np.where(c1 < c2, np.nan, c2)

        s_d = np.sqrt(c1 - c2)
        s_n = (h * np.cos(x) * np.cos(y) - s_d) / (np.cos(y) * np.cos(y) + fp * np.sin(y) * np.sin(y))
        s_1 = h - s_n * np.cos(x) * np.cos(y)
        s_2 = s_n * np.sin(x) * np.cos(y)
        s_3 = s_n * np.sin(y)
        s_xy = np.sqrt(s_1 * s_1 + s_2 * s_2)

        geographic_lon = np.arctan(s_2 / s_1) + self.sub_lon
        geographic_lat = np.arctan(fp * (s_3 / s_xy))

        geographic_lon *= RAD_TO_DEG
        geographic_lat *= RAD_TO_DEG

        return geographic_lon, geographic_lat

    def sat_to_earth(self, x, y):
        if self.scan_geom == 'GOES':
            x, y = goes_to_geos(x, y)

        c1 = (h * np.cos(x) * np.cos(y)) * (h * np.cos(x) * np.cos(y))
        c2 = (np.cos(y) * np.cos(y) + fp * np.sin(y) * np.sin(y)) * d

        if c1 < c2:
            return np.nan, np.nan

        s_d = np.sqrt(c1 - c2)
        s_n = (h * np.cos(x) * np.cos(y) - s_d) / (np.cos(y) * np.cos(y) + fp * np.sin(y) * np.sin(y))
        s_1 = h - s_n * np.cos(x) * np.cos(y)
        s_2 = s_n * np.sin(x) * np.cos(y)
        s_3 = s_n * np.sin(y)
        s_xy = np.sqrt(s_1 * s_1 + s_2 * s_2)

        geographic_lon = np.arctan(s_2 / s_1) + self.sub_lon
        geographic_lat = np.arctan(fp * (s_3 / s_xy))

        geographic_lon *= RAD_TO_DEG
        geographic_lat *= RAD_TO_DEG

        return geographic_lon, geographic_lat

    def check_inside(self, lon, lat):
        return True


def get_navigation(satellite='GOES16', domain='FD'):
    nav = None

    if satellite == 'GOES16':
        if domain == 'FD':
            nav = GEOSNavigation()
        elif domain == 'CONUS':
            nav = GEOSNavigation(sub_lon=-75.0, CFAC=5.6E-05, COFF=-0.101332,
                                 LFAC=-5.6E-05, LOFF=0.128212, num_elems=2500, num_lines=1500)
    elif satellite == 'GOES17':
        if domain == 'FD':
            nav = GEOSNavigation(sub_lon=-137.0)
        elif domain == 'CONUS':
            pass
    elif satellite == 'H08':
        nav = GEOSNavigation(sub_lon=140.7, barycentric_height=42164.0, scan_geom='GEOS',
                             CFAC=5.58879902955962e-05, LFAC=-5.58879902955962e-05,
                             COFF=-0.153719917308037, LOFF=0.153719917308037, num_elems=5500, num_lines=5500)
    elif satellite == 'H09':
        nav = GEOSNavigation(sub_lon=140.7, barycentric_height=42164.0, scan_geom='GEOS',
                             CFAC=5.58879902955962e-05, LFAC=-5.58879902955962e-05,
                             COFF=-0.153719917308037, LOFF=0.153719917308037, num_elems=5500, num_lines=5500)

    return nav


def get_lon_lat_2d_mesh(nav, ll, cc, offset=0):
    num_elems = len(cc)
    num_lines = len(ll)
    cc = np.array(cc)
    ll = np.array(ll)
    cc[:] += offset
    ll[:] += offset
    x_rad = cc * nav.CFAC + nav.COFF
    y_rad = ll * nav.LFAC + nav.LOFF

    ll, cc = np.meshgrid(ll, cc, indexing='ij')
    cc = cc.flatten()
    ll = ll.flatten()
    lon_s, lat_s = nav.lc_to_earth(cc, ll)
    lons_2d = lon_s.reshape((num_lines, num_elems))
    lats_2d = lat_s.reshape((num_lines, num_elems))

    return lons_2d, lats_2d, x_rad, y_rad

# 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)
#
#
# 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('-----------------------')