#
# 大气延迟校正工具计算
# 参考：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 osgeo import gdal,gdalconst
from tqdm import tqdm
from matplotlib import pyplot as plt

# 声明常数
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