#
# 大气延迟校正工具计算
# 参考：https://github.com/insarlab/PyAPS
#

import os
import numpy as np
import scipy.interpolate as si
import netCDF4 as Nc
import scipy.interpolate as intp
import scipy.integrate as intg
from tool.algorithm.image.ImageHandle import ImageHandler
from tqdm import tqdm

# 声明常数
constdict = {}
k1 = 0.776   #(K/Pa)
k2 = 0.716   #(K/Pa)
k3 = 3750    #(K^2.Pa)
g = 9.81     #(m/s^2)
Rd = 287.05  #(J/Kg/K)
Rv = 461.495 #(J/Kg/K)
mma = 29.97  #(g/mol)
mmH = 2.0158 #(g/mol)
mmO = 16.0   #(g/mol)
Rho = 1000.0 #(kg/m^3)

a1w = 611.21 # hPa
a3w = 17.502 #
a4w = 32.19  # K
a1i = 611.21 # hPa
a3i = 22.587 #
a4i = -0.7   # K
T3 = 273.16  # K
Ti = 250.16  # K
nhgt = 300   # Number of levels for interpolation (OLD 151)
minAlt = -200.0
maxAlt = 50000.0
minAltP = -200.0


def correct_single_ifgram(ifg_path, phas_path, cor_dis_file):
    data=ImageHandler.get_band_array(ifg_path,1)
    tropo=ImageHandler.get_band_array(phas_path,1)
    data =tropo - np.angle(data)
    return data
    
def createpha(master_pha,aux_pha,wavelen,out_pha_path):
    pha=master_pha-aux_pha
    #pha=pha*-4. * np.pi / float( wavelen)
    ImageHandler.write_imgArray(out_pha_path,pha)
    return pha

class ERA:
    """ 用于存储从ear 解析得到的参数"""
    def __init__(self) -> None:
        # 定义计算变量
        self.lon_array=None
        self.lat_array=None      
        self.level=None
        self.t=None
        self.r=None
        self.z=None
        pass

    def read_nc(self,nc_file,Extent=None,level_list=None):
        """
        function: 从气象数据中读取数组

        """
        try:
            nc_data_obj = Nc.Dataset(nc_file)
            inform = nc_data_obj.variables
            lon_array = inform['longitude'][:] # lon数组
            lat_array = inform['latitude'][:] # lat数组
            n_lat = lat_array.shape[0]             # 经线列数量
            n_lon = lon_array.shape[0]               # 纬线行数量
            level=np.array(inform["level"])         # pressure level in m 压力面
            t = np.array(inform["t"])               # temp--k  (1,37,latitude,longitude)
            r = np.array(inform["r"])               # relative_humidity--% (1,37,latitude,longitude)
            z = np.array(inform["z"])               # Geopotential--m**2 s**-2   (1,37,latitude,longitude)
            
            if not Extent is None:
                min_lat,max_lat,min_lon,max_lon=Extent
                mask_lon= (lon_array>=min_lon)&(lon_array<=max_lon)
                mask_lat= (lat_array>=min_lat)&(lat_array<=max_lat)
                mask_lon=np.where(mask_lon==1)[0]
                mask_lat=np.where(mask_lat==1)[0]
                lat_array=lat_array[mask_lat]
                lon_array=lat_array[mask_lon]
                t=t[:,:,mask_lat,:][:,:,:,mask_lon]
                r=r[:,:,mask_lat,:][:,:,:,mask_lon]
                z=z[:,:,mask_lat,:][:,:,:,mask_lon]
            if not level_list is None:
                mask_level=[]
                for l_item in level_list:
                    mask_level.append(np.where(level==l_item)[0][0])
                level=level[mask_level]
                t=t[:,mask_level,:,:]
                r=r[:,mask_level,:,:]
                z=z[:,mask_level,:,:]
            # 根据
            self.level=level
            self.lat_array=lat_array
            self.lon_array=lon_array
            self.r=r 
            self.t=t
            self.z=z
            # log_info("level:\n{}".format(str(self.level)))
            # log_info("lat_array:\n{}".format(str(self.lat_array)))
            # log_info("lon_array:\n{}".format(str(self.lon_array)))
            # log_info("r:\n{}".format(str(self.r)))
            # log_info("t:\n{}".format(str(self.t)))
            # log_info("z:\n{}".format(str(self.z)))  # (1,level,lat,lon)
        except Exception as e :
            print(e)
        #     log_err(e)
        # log_info("="*50)
        return True


class AtomZtdDelay:
    def __init__(self,atomdelay_path) -> None:
        self.init_var()
        self.const_var()
        self.output=atomdelay_path
        pass
    def const_var(self):
        self.lvls = np.array([1, 2, 3, 5, 7, 10, 20, 30, 50, 70, 100, 125,
                         150, 175,200, 225, 250, 300, 350, 400, 450, 
                         500, 550, 600, 650, 700, 750, 775, 800, 825, 
                         850, 875, 900, 925, 950, 975, 1000])
        self.nlvls = len(self.lvls)    
        self.hgtscale = ((maxAlt - minAlt) / nhgt) / 0.703
        self.bufspc = 1.2  
        self.alpha =  Rv / Rd
    
    def set_waveLen(self,wvl):
        self.wvl=wvl
    
    def set_incidenceAngle(self,inc_path):
        local_angle_array = ImageHandler.get_band_array(inc_path, 1)    # 获取局部入射角
        self.inc=local_angle_array
    
    def init_var(self):
        self.deley_all=None
        self.Dry_delay=None
        self.Wet_delay=None
        self.lon_arr=None
        self.lat_arr=None
        self.Pi=None
        self.Ti=None
        self.Vi=None
        self.hgt=None
        self.lvls=None
        self.gph=None
        self.tmp=None
        self.vpr=None
        self.dout=None
    
    def load_npy(self,output):
        # 加载数据
        # log_info("="*50)
        # log_info("加载预处理结果")
        # log_info("{}:{}".format(os.path.join(output,'deley_all.npy'),os.path.exists(os.path.join(output,'deley_all.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'Dry_delay.npy'),os.path.exists(os.path.join(output,'Dry_delay.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'Wet_delay.npy'),os.path.exists(os.path.join(output,'Wet_delay.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'lon_arr.npy'),os.path.exists(os.path.join(output,'lon_arr.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'lat_arr.npy'),os.path.exists(os.path.join(output,'lat_arr.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'Pi.npy'),os.path.exists(os.path.join(output,'Pi.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'Ti.npy'),os.path.exists(os.path.join(output,'Ti.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'Vi.npy'),os.path.exists(os.path.join(output,'Vi.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'hgt.npy'),os.path.exists(os.path.join(output,'hgt.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'lvls.npy'),os.path.exists(os.path.join(output,'lvls.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'gph.npy'),os.path.exists(os.path.join(output,'gph.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'tmp.npy'),os.path.exists(os.path.join(output,'tmp.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'vpr.npy'),os.path.exists(os.path.join(output,'vpr.npy'))))
        # log_info("{}:{}".format(os.path.join(output,'dout.npy'),os.path.exists(os.path.join(output,'dout.npy'))))
        # log_info("="*50)
        self.deley_all=np.load(os.path.join(output,'deley_all.npy'))
        self.Dry_delay=np.load(os.path.join(output,'Dry_delay.npy'))
        self.Wet_delay=np.load(os.path.join(output,'Wet_delay.npy'))
        self.lon_arr=np.load(os.path.join(output,'lon_arr.npy'))
        self.lat_arr=np.load(os.path.join(output,'lat_arr.npy'))
        self.Pi=np.load(os.path.join(output,'Pi.npy'))
        self.Ti=np.load(os.path.join(output,'Ti.npy'))
        self.Vi=np.load(os.path.join(output,'Vi.npy'))
        self.hgt=np.load(os.path.join(output,'hgt.npy'))
        self.lvls=np.load(os.path.join(output,'lvls.npy')).tolist()
        self.gph=np.load(os.path.join(output,'gph.npy'))
        self.tmp=np.load(os.path.join(output,'tmp.npy'))
        self.vpr=np.load(os.path.join(output,'vpr.npy'))
        self.dout=np.load(os.path.join(output,'dlosout.npy'))
        self.pha=np.load(os.path.join(output,'pha.npy'))
    
    def save(self):
        # 保存成临时结果
        ImageHandler.write_imgArray(os.path.join(self.output,'deley_all.tiff'),self.deley_all)
        # np.save(os.path.join(self.output,'Dry_delay.npy'),self.Dry_delay)
        # np.save(os.path.join(self.output,'Wet_delay.npy'),self.Wet_delay)
        # np.save(os.path.join(self.output,'lon_arr.npy'),np.array(self.lon_arr))
        # np.save(os.path.join(self.output,'lat_arr.npy'),np.array(self.lat_arr))
        # np.save(os.path.join(self.output,'Pi.npy'),self.Pi)
        # np.save(os.path.join(self.output,'Ti.npy'),self.Ti)
        # np.save(os.path.join(self.output,'Vi.npy'),self.Vi)
        # np.save(os.path.join(self.output,'hgt.npy'),self.hgt)
        # np.save(os.path.join(self.output,'lvls.npy'),np.array(self.lvls))
        # np.save(os.path.join(self.output,'gph.npy'),self.gph)
        # np.save(os.path.join(self.output,'tmp.npy'),self.tmp)
        # np.save(os.path.join(self.output,'vpr.npy'),self.vpr)
        ImageHandler.write_imgArray(os.path.join(self.output,'dlosout.tiff'),self.dout)  # ztd
        ImageHandler.write_imgArray(os.path.join(self.output,'pha.tiff'),self.pha)

    def get_dem(self,dem_path):
        dem_array=ImageHandler.get_band_array(dem_path,1)
        self.dem=dem_array
        return True
    
    def get_lon_lat(self,reflect_lon_path,reflect_lat_path):
        lon_table=ImageHandler.get_band_array(reflect_lon_path,1)
        lat_table=ImageHandler.get_band_array(reflect_lat_path,1)   
        
        self.lon=lon_table
        self.lat=lat_table
        return True
     
    def get_delay(self):
        # log_info("计算相位")
        minAltp=-200.0
        
        lon_arr=self.lon
        lat_arr=self.lat
        # 根据行列号
        min_lon=lon_arr.min()-self.bufspc
        max_lon=lon_arr.max()+self.bufspc
        min_lat=lat_arr.min()-self.bufspc
        max_lat=lat_arr.max()+self.bufspc
        # WGS  [-180,-90,180,90] 
        # dem 
        shp=lon_arr.shape
        lon_arr=lon_arr.reshape(-1)
        if min_lon<0 and max_lon>0:
            lon_arr[np.where(lon_arr<0)]=lon_arr[np.where(lon_arr<0)]+360
        lon_arr=lon_arr.reshape(shp)
        self.ny,self.nx=self.dem.shape
        
        dout = np.zeros((self.ny, self.nx), dtype=np.float32)
        
        intp_1d = si.interp1d(self.hgt, self.deley_all, kind='cubic', axis=-1)
        
        self.dem[np.isnan(self.dem)] = minAltp
        self.dem[self.dem < minAltp] = minAltp
        self.mask=np.ones(self.dem.shape)
        minH = np.max([np.nanmin(self.dem), self.hgt.min()])
        maxH = int(np.nanmax(self.dem)) + 100.
        kh = np.arange(minH,maxH)
        self.Delfn_1m = intp_1d(kh)
        self.alti = kh        
        # no reshape
        
        Lonu = self.lon_arr[0,:]
        Latu = self.lat_arr[:,0]
        linearint = si.RegularGridInterpolator(
            (Latu[::-1], Lonu,kh),
            self.Delfn_1m[::-1,:,:],
            method='linear',
            bounds_error=False,
            fill_value=0.0,
        )
        wvl=self.wvl
        # Loop on the lines
        pha=dout.copy()
        # log_info("计算延迟值")
        for m in tqdm(range(self.ny)):

            # Get latitude and longitude arrays
            lati = self.lat[m,:]*self.mask[m,:]
            loni = self.lon[m,:]*self.mask[m,:]

            # Remove negative values

            # Remove NaN values
            ii = np.where(np.isnan(lati))
            jj = np.where(np.isnan(loni))
            xx = np.union1d(ii,jj)
            lati[xx]=0.0
            loni[xx]=0.0

            cinc = np.cos(self.inc[m,:]*np.pi/180.)

            # Make the bilinear interpolation
            D = self.dem[m,:]
            val = linearint(np.vstack((lati, loni, D)).T)
            val[xx] = np.nan

            # save output
            dout[m,:] = val
            pha[m,:]=val*np.pi*4.0/(cinc*wvl)
            
        self.dout=dout
        self.pha= -1*pha
        return pha
        pass
    
    def get_deley_dlf(self,nc_file):
        # log_info("计算延迟路径")
        [lon_array,lat_array, lvls, gph,tmp,vpr]=self.get_ear5_nc(nc_file)
        lon_array,lat_array=np.meshgrid(lon_array,lat_array)
        hgt = np.linspace(minAltP, gph.max().round(), nhgt)
        [Pi,Ti,Vi]=self.intP2H(lvls,hgt,gph,tmp,vpr)
        [Dry_delay,Wet_delay] = self.PTV2del(Pi,Ti,Vi,hgt)
        deley_all = Dry_delay+Wet_delay
        # 保存结果
        self.deley_all=deley_all
        self.Dry_delay=Dry_delay
        self.Wet_delay=Wet_delay
        self.lon_arr=lon_array
        self.lat_arr=lat_array
        self.Pi=Pi
        self.Ti=Ti
        self.Vi=Vi
        self.hgt=hgt
        self.lvls=lvls
        self.gph=gph # 高度
        self.tmp=tmp # 温度
        self.vpr=vpr # 湿度
        return True
    
    def get_ear5_nc(self,nc_file,Extent=None):
        era= ERA()
        era.read_nc(nc_file,Extent,self.lvls)
        lvls=self.lvls*100 # 转为 绝对压力
        nlvls=len(era.level)
        nlat=len(era.lat_array)
        nlon=len(era.lon_array)
        # 逐层分析
        gph=1.0*era.z[0,:,:,:]/g
        tmp=era.t[0,:,:,:]
        vpr=era.r[0,:,:,:]
        for i in range(nlvls):
            esat = self.cc_era(tmp[i,:,:])
            vpr[i,:,:]=(vpr[i,:,:]/100.0)*esat
        # 
        return era.lon_array,era.lat_array, lvls, gph,tmp,vpr

    def intP2H(self,lvls,hgt,gph,tmp,vpr):
        # Interpolate pressure, temperature and Humidity of hgt
        #minAlt = minAlt     #Hardcoded parameter. -200.0
        maxAlt_2 = gph.max().round()
        
        nlat = gph.shape[1]           #Number of stations
        nlon = gph.shape[2]
        nhgt = len(hgt)               #Number of height points
        Presi = np.zeros((nlat,nlon,nhgt))
        Tempi = np.zeros((nlat,nlon,nhgt))
        Vpri  = np.zeros((nlat,nlon,nhgt))

        for i in range(nlat):
            for j in range(nlon):
                temp = gph[:,i,j]       #Obtaining height values
                hx = temp.copy()
                sFlag = False
                eFlag = False
                #Add point at start
                if (hx.min() > minAlt): # -200.0
                    sFlag = True
                    #changed from 1 to 0 (-1 should also work), CL
                    hx = np.concatenate((hx,[minAlt-1]),axis = 0)

                #Add point at end
                if (hx.max() < maxAlt_2):
                    eFlag = True
                    #changed from 1 to 0 (-1 should also work), CL
                    hx = np.concatenate(([maxAlt_2+1],hx),axis=0)

                #Splines needs monotonically increasing.
                hx = -hx

                #Interpolating pressure values
                hy = lvls.copy()
                if (sFlag == True):
                    val = hy[-1] +(hx[-1] - hx[-2])* (hy[-1] - hy[-2])/(hx[-2]-hx[-3])
                    #changed from 1 to 0 (-1 should also work), CL
                    hy = np.concatenate((hy,[val]),axis=0)

                if (eFlag == True):
                    val = hy[0] - (hx[0] - hx[1]) * (hy[0] - hy[1])/(hx[1]-hx[2])
                    #changed from 1 to 0 (-1 should also work), CL
                    hy = np.concatenate(([val],hy),axis=0)

                tck = intp.interp1d(hx,hy,kind='cubic')
                #Again negative for consistency with hx
                temp = tck(-hgt)
                Presi[i,j,:] = temp.copy()
                del temp

                #Interpolating temperature
                temp = tmp[:,i,j]
                hy = temp.copy()
                if (sFlag == True):
                    val = hy[-1] +(hx[-1] - hx[-2])* (hy[-1] - hy[-2])/(hx[-2]-hx[-3])
                    #changed from 1 to 0 (-1 should also work), CL
                    hy = np.concatenate((hy,[val]),axis=0)

                if (eFlag == True):
                    val = hy[0] - (hx[0] - hx[1]) * (hy[0] - hy[1])/(hx[1]-hx[2])
                    #changed from 1 to 0 (-1 should also work), CL
                    hy = np.concatenate(([val],hy),axis=0)

                tck = intp.interp1d(hx,hy,kind='cubic')
                temp = tck(-hgt)
                Tempi[i,j,:] = temp.copy()
                del temp

                #Interpolating vapor pressure
                temp = vpr[:,i,j]
                hy = temp.copy()
                if (sFlag == True):
                    val = hy[-1] +(hx[-1] - hx[-2])* (hy[-1] - hy[-2])/(hx[-2]-hx[-3])
                    #changed from 1 to 0 (-1 should also work), CL
                    hy = np.concatenate((hy,[val]),axis=0)

                if (eFlag == True):
                    val = hy[0] - (hx[0] - hx[1]) * (hy[0] - hy[1])/(hx[1]-hx[2])
                    #changed from 1 to 0 (-1 should also work), CL
                    hy = np.concatenate(([val],hy),axis=0)

                tck = intp.interp1d(hx,hy,kind='cubic')
                temp = tck(-hgt)
                Vpri[i,j,:] = temp.copy()
                del temp

        return Presi,Tempi,Vpri    
    
    def PTV2del(self,Presi,Tempi,Vpri,hgt):
        '''Computes the delay function given Pressure, Temperature and Vapor pressure.

        Args:
            * Presi (np.ndarray) : Pressure at height levels.
            * Tempi (np.ndarray) : Temperature at height levels.
            * Vpri  (np.ndarray) : Vapor pressure at height levels.
            * hgt   (np.ndarray) : Height levels.
            * cdict (np.ndarray) : Dictionary of constants.

        Returns:
            * DDry2 (np.ndarray) : Dry component of atmospheric delay.
            * DWet2 (np.ndarray) : Wet component of atmospheric delay.

        .. note::
            Computes refractive index at each altitude and integrates the delay using cumtrapz.'''
        nhgt = len(hgt)              #Number of height points
        nlat = Presi.shape[0]        #Number of stations
        nlon = Presi.shape[1]
        WonT = Vpri/Tempi
        WonT2 = WonT/Tempi


        #Dry delay
        DDry2 = np.zeros((nlat,nlon,nhgt))
        DDry2[:,:,:] = k1*Rd*(Presi[:,:,:] - Presi[:,:,-1][:,:,np.newaxis])*1.0e-6/g

        #Wet delay
        S1 = intg.cumtrapz(WonT,x=hgt,axis=-1)
        val = 2*S1[:,:,-1]-S1[:,:,-2]
        val = val[:,:,None]
        S1 = np.concatenate((S1,val),axis=-1)
        del WonT

        S2 = intg.cumtrapz(WonT2,x=hgt,axis=-1)
        val = 2*S2[:,:,-1]-S2[:,:,-2]
        val = val[:,:,None]
        S2 = np.concatenate((S2,val),axis=-1)
        DWet2 = -1.0e-6*((k2-k1*Rd/Rv)*S1+k3*S2)
        
        for i in range(nlat):
            for j in range(nlon):
                DWet2[i,j,:]  = DWet2[i,j,:] - DWet2[i,j,-1]

        return DDry2,DWet2    
        
    def cc_era(self,tmp):
        shp=tmp.shape
        tmp=tmp.reshape(-1)
        esatw = a1w*np.exp(a3w*(tmp-T3)/(tmp-a4w))
        esati = a1i*np.exp(a3i*(tmp-T3)/(tmp-a4i))
        result = np.zeros(esati.shape)
        mask_T3=tmp>=T3
        mask_Ti=tmp<=Ti
        mask_wgt=(mask_T3+mask_Ti)==0
        mask_T3=np.where(mask_T3==1)
        mask_Ti=np.where(mask_Ti==1)
        mask_wgt=np.where(mask_wgt==1)
        result[mask_T3]=esatw[mask_T3]
        result[mask_Ti]=esati[mask_Ti]
        result[mask_wgt]=esati[mask_wgt]+(esatw[mask_wgt]-esati[mask_wgt])*np.square((tmp[mask_wgt]-Ti)/(T3-Ti)) 
        result=result.reshape(shp)
        return result