"""
Created on Thu Sep 25 20:02:25 2025

@author: my friend cjj
"""

import numpy as np
import math as mt
import netCDF4 as nc
import xarray as xr

import matplotlib.pyplot as plt
import matplotlib.ticker as mtick

import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.util import add_cyclic_point
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
from cartopy.mpl import ticker

from eofs.multivariate.standard import MultivariateEof
from eofs.standard import Eof

from scipy.signal import detrend
from scipy.stats import linregress

import metpy.plots as mp
from metpy.units import units


plt.rcParams['font.size'] = 15  # 字体大小
plt.rcParams['font.family'] = 'Times New Roman'  # 字体类型，支持中文显示
plt.rcParams['font.weight'] = 'normal'  # 字体粗细


#作业1-3
path='D:/PHD/lesson_data/HadISST_sst.nc'
da1=xr.open_dataset(path)
#print(da1)

#转-180-180经度数据为0-360，即将-180-0改为180-360
da1['lon'] = xr.where(da1['longitude'] <0 , da1['longitude'] + 360, da1['longitude'])
da1 = (da1
      .swap_dims({'longitude': 'lon'})
      .sel(**{'lon': sorted(da1.lon)})
      .drop('longitude'))
da1 = da1.rename({'lon': 'longitude'})
print(da1)

#SST=da1.sel(latitude=slice(30,-30),longitude=slice(100,300),time=slice("1981","2024"))
leftlon,rightlon,lowerlat,upperlat=(110,250,20,70)#设置经纬度范围
SST=da1.sel(longitude=slice(leftlon,rightlon),latitude=slice(upperlat,lowerlat),time=slice("1900-06","2025-02"))
sst=np.array(SST.sst[:])
lat=np.array(SST.latitude[:])
lon=np.array(SST.longitude[:])
time=np.array(SST.time[:])

sst = np.where(sst == -1000, np.nan, sst)

sst_2=(sst[0:1505:12,:,:]+sst[1:1505:12,:,:]+sst[2::12,:,:])/3#夏季平均,1979至2024年夏季
#距平
sst_jja=sst_2-np.mean(sst_2,axis=0)


# 3. 二维化、权重、掩膜
nt, ny, nx = sst_jja.shape
X = sst_jja.reshape(nt, -1)
mask = ~np.isnan(X).any(axis=0)
X_clean = X[:, mask] 

# 4. 去趋势
X_clean = detrend(X_clean, axis=0, type='linear')

#直接回填
td2=np.full((nt,ny*nx), np.nan)
td2[:,mask]=X_clean
td2=td2.reshape([nt,ny,nx])

#计算纬度权重
coslat=np.cos(np.deg2rad(lat))
wgts = np.sqrt(coslat)[..., np.newaxis]
#创建EOF分解器
solver=Eof(td2,weights=wgts)
eof=solver.eofsAsCorrelation(neofs=4)
#此处的neofs的值是我们需要的空间模态数，比如这里我们打算画四个模态
pc=solver.pcs(npcs=4,pcscaling=1)#方差
var=solver.varianceFraction(neigs=4)


fig=plt.figure(figsize=(9,9),dpi=200)#设置画布
proj=ccrs.PlateCarree(central_longitude=180)
lon_formatter=ticker.LongitudeFormatter()
lat_formatter=ticker.LatitudeFormatter()


fig_ax1=fig.add_axes([0.1,0.95,0.5,0.3],projection=proj)
fig_ax1.set_extent([leftlon,rightlon,lowerlat,upperlat],crs=ccrs.PlateCarree())
fig_ax1.add_feature(cfeature.OCEAN,edgecolor='black')
fig_ax1.add_feature(cfeature.LAND, color='lightgrey')
fig_ax1.add_feature(cfeature.COASTLINE,lw=1)
fig_ax1.set_xticks(np.arange(leftlon+10,rightlon,20),crs=ccrs.PlateCarree())
fig_ax1.set_yticks(np.arange(lowerlat,upperlat+5,10),crs=ccrs.PlateCarree())
fig_ax1.xaxis.set_major_formatter(lon_formatter)
fig_ax1.yaxis.set_major_formatter(lat_formatter)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
rivers_110m = cfeature.NaturalEarthFeature('physical', 'rivers_lake_centerlines', '110m')
fig_ax1.set_title('(a) summer pattern1',loc='left',fontsize =15)
fig_ax1.set_title( '%.2f%%' % (var[0]*100),loc='right',fontsize =15)
c1=fig_ax1.contourf(lon,lat, eof[0,:,:], levels=np.arange(-0.9,0.9,0.1),zorder=0, extend = 'both',transform=ccrs.PlateCarree(), cmap=plt.cm.RdBu_r)

fig_ax1=fig.add_axes([0.1,0.7,0.5,0.3],projection=proj)
fig_ax1.set_extent([leftlon,rightlon,lowerlat,upperlat],crs=ccrs.PlateCarree())
fig_ax1.add_feature(cfeature.OCEAN,edgecolor='black')
fig_ax1.add_feature(cfeature.LAND, color='lightgrey')
fig_ax1.add_feature(cfeature.COASTLINE,lw=1)
fig_ax1.set_xticks(np.arange(leftlon+10,rightlon,20),crs=ccrs.PlateCarree())
fig_ax1.set_yticks(np.arange(lowerlat,upperlat+5,10),crs=ccrs.PlateCarree())
fig_ax1.xaxis.set_major_formatter(lon_formatter)
fig_ax1.yaxis.set_major_formatter(lat_formatter)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
rivers_110m = cfeature.NaturalEarthFeature('physical', 'rivers_lake_centerlines', '110m')
fig_ax1.set_title('(b) summer pattern2',loc='left',fontsize =15)
fig_ax1.set_title( '%.2f%%' % (var[1]*100),loc='right',fontsize =15)
c1=fig_ax1.contourf(lon,lat, eof[1,:,:],levels=np.arange(-0.9,0.9,0.1), zorder=0, extend = 'both',transform=ccrs.PlateCarree(), cmap=plt.cm.RdBu_r)

t=np.arange(1900,2025)
fig_ax5=fig.add_axes([0.65,1.01,0.4,0.18])
fig_ax5.set_title('(e) summer PC1',loc='left',fontsize = 15)
fig_ax5.set_ylim(-3.5,3.5)
fig_ax5.axhline(0,linestyle="--")
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.xlim(1900,2024)
fig_ax5.plot(t,pc[:,0],color='k',linewidth=1.5,label='PC')


fig_ax5=fig.add_axes([0.65,0.76,0.4,0.18])
fig_ax5.set_title('(f) summer PC2',loc='left',fontsize = 15)
fig_ax5.set_ylim(-3.5,3.5)
fig_ax5.axhline(0,linestyle="--")
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.xlim(1900,2024)
fig_ax5.plot(t,pc[:,1],color='k',linewidth=1.5,label='PC')

cbposition=fig.add_axes([0.1, 0.69, 0.95, 0.02])
cb=fig.colorbar(c1,cax=cbposition,orientation='horizontal',format='%.1f')
cb.ax.tick_params(labelsize=14)
#cb.set_ticks([-0.8,-0.4,0,0.4,0.8])

#直接改成冬季
sst_2=(sst[6:1505:12,:,:]+sst[7:1505:12,:,:]+sst[8::12,:,:])/3
#距平
sst_jja=sst_2-np.mean(sst_2,axis=0)


# 3. 二维化、权重、掩膜
nt, ny, nx = sst_jja.shape
X = sst_jja.reshape(nt, -1)
mask = ~np.isnan(X).any(axis=0)
X_clean = X[:, mask] 

# 4. 去趋势
X_clean = detrend(X_clean, axis=0, type='linear')

#直接回填
td2=np.full((nt,ny*nx), np.nan)
td2[:,mask]=X_clean
td2=td2.reshape([nt,ny,nx])

#计算纬度权重
coslat=np.cos(np.deg2rad(lat))
wgts = np.sqrt(coslat)[..., np.newaxis]
#创建EOF分解器
solver=Eof(td2,weights=wgts)
eof=solver.eofsAsCorrelation(neofs=4)
#此处的neofs的值是我们需要的空间模态数，比如这里我们打算画四个模态
pc=solver.pcs(npcs=4,pcscaling=1)#方差
var=solver.varianceFraction(neigs=4)


fig=plt.figure(figsize=(9,9),dpi=200)#设置画布
proj=ccrs.PlateCarree(central_longitude=180)
lon_formatter=ticker.LongitudeFormatter()
lat_formatter=ticker.LatitudeFormatter()


fig_ax1=fig.add_axes([0.1,0.95,0.5,0.3],projection=proj)
fig_ax1.set_extent([leftlon,rightlon,lowerlat,upperlat],crs=ccrs.PlateCarree())
fig_ax1.add_feature(cfeature.OCEAN,edgecolor='black')
fig_ax1.add_feature(cfeature.LAND, color='lightgrey')
fig_ax1.add_feature(cfeature.COASTLINE,lw=1)
fig_ax1.set_xticks(np.arange(leftlon+10,rightlon,20),crs=ccrs.PlateCarree())
fig_ax1.set_yticks(np.arange(lowerlat,upperlat+5,10),crs=ccrs.PlateCarree())
fig_ax1.xaxis.set_major_formatter(lon_formatter)
fig_ax1.yaxis.set_major_formatter(lat_formatter)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
rivers_110m = cfeature.NaturalEarthFeature('physical', 'rivers_lake_centerlines', '110m')
fig_ax1.set_title('(a) winter pattern1',loc='left',fontsize =15)
fig_ax1.set_title( '%.2f%%' % (var[0]*100),loc='right',fontsize =15)
c1=fig_ax1.contourf(lon,lat, eof[0,:,:], levels=np.arange(-0.9,0.9,0.1),zorder=0, extend = 'both',transform=ccrs.PlateCarree(), cmap=plt.cm.RdBu_r)

fig_ax1=fig.add_axes([0.1,0.7,0.5,0.3],projection=proj)
fig_ax1.set_extent([leftlon,rightlon,lowerlat,upperlat],crs=ccrs.PlateCarree())
fig_ax1.add_feature(cfeature.OCEAN,edgecolor='black')
fig_ax1.add_feature(cfeature.LAND, color='lightgrey')
fig_ax1.add_feature(cfeature.COASTLINE,lw=1)
fig_ax1.set_xticks(np.arange(leftlon+10,rightlon,20),crs=ccrs.PlateCarree())
fig_ax1.set_yticks(np.arange(lowerlat,upperlat+5,10),crs=ccrs.PlateCarree())
fig_ax1.xaxis.set_major_formatter(lon_formatter)
fig_ax1.yaxis.set_major_formatter(lat_formatter)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
rivers_110m = cfeature.NaturalEarthFeature('physical', 'rivers_lake_centerlines', '110m')
fig_ax1.set_title('(b) winter pattern2',loc='left',fontsize =15)
fig_ax1.set_title( '%.2f%%' % (var[1]*100),loc='right',fontsize =15)
c1=fig_ax1.contourf(lon,lat, eof[1,:,:],levels=np.arange(-0.9,0.9,0.1), zorder=0, extend = 'both',transform=ccrs.PlateCarree(), cmap=plt.cm.RdBu_r)

t=np.arange(1900,2025)
fig_ax5=fig.add_axes([0.65,1.01,0.4,0.18])
fig_ax5.set_title('(e) summer PC1',loc='left',fontsize = 15)
fig_ax5.set_ylim(-3.5,3.5)
fig_ax5.axhline(0,linestyle="--")
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.xlim(1900,2024)
fig_ax5.plot(t,pc[:,0],color='k',linewidth=1.5,label='PC')


fig_ax5=fig.add_axes([0.65,0.76,0.4,0.18])
fig_ax5.set_title('(f) summer PC2',loc='left',fontsize = 15)
fig_ax5.set_ylim(-3.5,3.5)
fig_ax5.axhline(0,linestyle="--")
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.xlim(1900,2024)
fig_ax5.plot(t,pc[:,1],color='k',linewidth=1.5,label='PC')

cbposition=fig.add_axes([0.1, 0.69, 0.95, 0.02])
cb=fig.colorbar(c1,cax=cbposition,orientation='horizontal',format='%.1f')
cb.ax.tick_params(labelsize=14)
#cb.set_ticks([-0.8,-0.4,0,0.4,0.8])


#作业2
path1='D:/PHD/lesson_data/NA_Tforecast/air.mon.mean.nc'
df1=xr.open_dataset(path1)
#df=df.sel(time=slice("2012","2025"),lon=slice(25,170),lat=slice(80,-10),level=300)
#leftlon,rightlon,lowerlat,upperlat=(70,150,5,55)#设置经纬度范围
#df2=df.sel(time=slice("1978","2025"),lon=slice(leftlon,rightlon),lat=slice(upperlat,lowerlat),level=1000)
lat1=np.array(df1['lat'])
lon1=np.array(df1['lon'])
level=np.array(df1['level'])
print(df1)
air=np.array(df1.air)

#纬向平均
air=np.nanmean(air,axis=3)
#春夏年份+1
year1=78
year2=77
zz=17#17层高度
latt=73
t1=np.arange(year1)
t2=np.arange(year2)
#分季节，春夏秋冬
airs1=(air[2::12,:,:]+air[3::12,:,:]+air[4::12,:,:])/3
airs2=(air[5::12,:,:]+air[6::12,:,:]+air[7::12,:,:])/3
airs3=(air[8::12,:,:]+air[9::12,:,:]+air[10::12,:,:])/3
airs4=(air[11::12,:,:]+air[12::12,:,:]+air[13::12,:,:])/3
#tend
trend=np.zeros([4,zz,latt])
for i in range(zz):
    for j in range(latt):
        slope, intercept, r_value, p_value, std_err = linregress(t1, airs1[:,i,j])
        trend[0,i,j]=slope
        slope, intercept, r_value, p_value, std_err = linregress(t1, airs2[:,i,j])
        trend[1,i,j]=slope
        slope, intercept, r_value, p_value, std_err = linregress(t2, airs3[:,i,j])
        trend[2,i,j]=slope
        slope, intercept, r_value, p_value, std_err = linregress(t2, airs4[:,i,j])
        trend[3,i,j]=slope
                
#绘制高度-经度剖面图
fig=plt.figure(figsize=(9,9),dpi=200)#设置画布
fig_ax1=fig.add_axes([0.1,0.95,0.5,0.4])
a=np.linspace(-0.2,0.2,21)
b=np.linspace(-0.2,0.2,21)
c=fig_ax1.contourf(lat1,level, trend[0,:,:], levels=a,zorder=0, extend = 'both',cmap=plt.cm.RdBu_r ) 
contour = fig_ax1.contour(lat1,level,trend[0,:,:], levels=b,colors='black', linewidths=0.5)
fig_ax1.clabel(contour, inline=True, fontsize=12)
fig.gca().invert_yaxis()
wanted_ticks = [10,20,30,50, 100, 200, 300, 500,700, 850, 1000]
plt.yscale('log') 
fig_ax1.set_yticks(wanted_ticks)

# 自动给 y 轴加“hPa”
fig_ax1.yaxis.set_major_formatter(mtick.FuncFormatter(lambda v, pos: f'{int(v)} hPa'))
fig_ax1.xaxis.set_major_locator(plt.MaxNLocator(7))  # 控制x轴刻度的数量
fig_ax1.xaxis.set_major_formatter(LatitudeFormatter())
#cb=plt.colorbar(c,orientation='horizontal',format='%.2f',shrink=1,pad=0.1)
plt.title('(a) spring trend zonal-mean',loc='left')


fig_ax1=fig.add_axes([0.7,0.95,0.5,0.4])
c=fig_ax1.contourf(lat1,level, trend[1,:,:], levels=a,zorder=0, extend = 'both',cmap=plt.cm.RdBu_r ) 
contour = fig_ax1.contour(lat1,level,trend[1,:,:], levels=b,colors='black', linewidths=0.5)
fig_ax1.clabel(contour, inline=True,fontsize=12)
fig.gca().invert_yaxis()
plt.yscale('log') 
fig_ax1.set_yticks(wanted_ticks)
fig_ax1.yaxis.set_major_formatter(mtick.FuncFormatter(lambda v, pos: f'{int(v)} hPa'))# 自动给 y 轴加“hPa”
fig_ax1.xaxis.set_major_locator(plt.MaxNLocator(7))  # 控制x轴刻度的数量
fig_ax1.xaxis.set_major_formatter(LatitudeFormatter())
#cb=plt.colorbar(c,orientation='horizontal',format='%.2f',shrink=1,pad=0.1)
plt.title('(b) summer trend zonal-mean',loc='left')

fig_ax1=fig.add_axes([0.1,0.48,0.5,0.4])
c=fig_ax1.contourf(lat1,level, trend[2,:,:], levels=a,zorder=0, extend = 'both',cmap=plt.cm.RdBu_r ) 
contour = fig_ax1.contour(lat1,level,trend[2,:,:], levels=b,colors='black', linewidths=0.5)
fig_ax1.clabel(contour, inline=True,fontsize=12)
fig.gca().invert_yaxis()
plt.yscale('log') 
fig_ax1.set_yticks(wanted_ticks)

fig_ax1.yaxis.set_major_formatter(mtick.FuncFormatter(lambda v, pos: f'{int(v)} hPa'))# 自动给 y 轴加“hPa”
fig_ax1.xaxis.set_major_locator(plt.MaxNLocator(7))  # 控制x轴刻度的数量
fig_ax1.xaxis.set_major_formatter(LatitudeFormatter())
#cb=plt.colorbar(c,orientation='horizontal',format='%.2f',shrink=1,pad=0.1)
plt.title('(c) autumn trend zonal-mean',loc='left')

fig_ax1=fig.add_axes([0.7,0.48,0.5,0.4])
c=fig_ax1.contourf(lat1,level, trend[3,:,:], levels=a,zorder=0, extend = 'both',cmap=plt.cm.RdBu_r ) 
contour = fig_ax1.contour(lat1,level,trend[3,:,:], levels=b,colors='black', linewidths=0.5)
fig_ax1.clabel(contour, inline=True, fontsize=12)
fig.gca().invert_yaxis()
plt.yscale('log') 
fig_ax1.set_yticks(wanted_ticks)
fig_ax1.yaxis.set_major_formatter(mtick.FuncFormatter(lambda v, pos: f'{int(v)} hPa'))# 自动给 y 轴加“hPa”
fig_ax1.xaxis.set_major_locator(plt.MaxNLocator(7))  # 控制x轴刻度的数量
fig_ax1.xaxis.set_major_formatter(LatitudeFormatter())
#cb=plt.colorbar(c,orientation='horizontal',format='%.2f',shrink=1,pad=0.1)
plt.title('(d) winter trend zonal-mean',loc='left')

cbar_ax = fig.add_axes([0.12, 0.4, 1, 0.025])
cb=fig.colorbar(c,cax=cbar_ax,orientation='horizontal',format='%.3f')