import numpy as 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
scan_geom = 'GOES'
# GOES-17 sub_lon = -137.0
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


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)

        if r_1 > 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)

        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)

        np.where(r_1 > self.h, lamda_sat == np.nan, lamda_sat)
        np.where(r_1 > self.h, theta_sat == 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)
        np.where(np.isnan(cc), cc == -1, cc)
        np.where(np.isnan(ll), ll == -1, ll)
        return cc, ll

    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 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':
            pass
    elif satellite == 'GOES17':
        if domain == 'FD':
            nav = GEOSNavigation(sub_lon=-137.0)
        elif domain == 'CONUS':
            pass

    return nav