for cube in [cube1, cube2]:
    for coord_name in [um.z, um.y, um.x]:
        cube.coord(coord_name).var_name = coord_name  # or None

import iris
import warnings
import numpy as np
import matplotlib.pyplot as plt
import aeolus

import iris
import warnings
import numpy as np
import matplotlib.pyplot as plt
import aeolus

from iris.analysis.cartography import _meshgrid, rotate_pole, rotate_winds

from aeolus.coord import get_xy_coords
from aeolus.calc.tl import regrid_to_tidally_locked_coordinates, regrid_to_rotated_pole_coordinates, rotate_winds_to_tidally_locked_coordinates
from aeolus.calc.stats import spatial
from aeolus.const import init_const
pcb_const=init_const('proxb')
warnings.filterwarnings("ignore")

pcb_2d = iris.load('pcb_winds_2d.nc')

def regrid_cubelist(cubes):            
    for cube in cubes:
        if cube.standard_name == 'eastward_wind':
            x_wind = cube[:,:].copy()
        if cube.standard_name == 'northward_wind':
            y_wind = cube[:,:].copy()
        if cube.standard_name == 'surface_temperature':
            tsurf = cube[:,:].copy()
        if cube.standard_name =='air_pressure':
            pressure = cube[:,:].copy()
    x_wind = x_wind.regrid(y_wind, iris.analysis.Linear())
    y_wind.coord('latitude').coord_system = iris.coord_systems.GeogCS(semi_major_axis= pcb_const.radius.data, longitude_of_prime_meridian=0)    
    y_wind.coord('longitude').coord_system = iris.coord_systems.GeogCS(semi_major_axis= pcb_const.radius.data, longitude_of_prime_meridian=0)    
    x_wind.coord('latitude').coord_system = iris.coord_systems.GeogCS(semi_major_axis= pcb_const.radius.data, longitude_of_prime_meridian=0)
    x_wind.coord('longitude').coord_system = iris.coord_systems.GeogCS(semi_major_axis= pcb_const.radius.data, longitude_of_prime_meridian=0)
    tsurf.coord('latitude').coord_system = iris.coord_systems.GeogCS(semi_major_axis= pcb_const.radius.data, longitude_of_prime_meridian=0)  
    tsurf.coord('longitude').coord_system = iris.coord_systems.GeogCS(semi_major_axis= pcb_const.radius.data, longitude_of_prime_meridian=0)

    x_wind, y_wind = rotate_winds_to_tidally_locked_coordinates(x_wind, y_wind, pole_lon=0, pole_lat=0)
    
    y_wind=regrid_to_tidally_locked_coordinates(y_wind)
    x_wind=regrid_to_tidally_locked_coordinates(x_wind)
    tsurf=regrid_to_tidally_locked_coordinates(tsurf)

    new_cubes = iris.cube.CubeList([tsurf, x_wind, y_wind])
    new_cubes[0].standard_name = 'surface_temperature'
    new_cubes[1].standard_name = 'eastward_wind'
    new_cubes[2].standard_name = 'northward_wind'
    return new_cubes

pcb_2d_regrid=regrid_cubelist(pcb_2d)

def plot_tsurf(cubes, time_mean=False, tl_coord=False):            
    for cube in cubes:
        if cube.standard_name == 'eastward_wind':
            x_wind = cube[:,:].copy()
        if cube.standard_name == 'northward_wind':
            y_wind = cube[:,:].copy()
        if cube.standard_name =='surface_temperature':
            temperature = cube[:,:].copy()
    
    y_wind = y_wind.regrid(x_wind, iris.analysis.Linear())
    xlon = x_wind.coord('longitude')
    ylon = y_wind.coord('longitude')
    print('Minimum temp:', spatial(temperature, 'min').data)
    print('Max temp:', spatial(temperature, 'max').data)
    
    xe = xlon.points
    ye = y_wind.coord('latitude').points
    ue = x_wind[:, :].data
    ve = y_wind[:, :].data
    #print(temperature[25,:,:].coord(''))

    fig = plt.figure(figsize=(10,8))
    ax = fig.add_subplot(111)
    c=plt.contourf(temperature[:,:].coord('longitude').points,temperature[:,:].coord('latitude').points,
                   temperature[:,:].data, levels=25, extend='both',cmap='RdBu_r')
    c=plt.contourf(temperature[:,:].coord('longitude').points,temperature[:,:].coord('latitude').points,
                   temperature[:,:].data, levels=40, extend='both',cmap='RdBu_r')
    cbar = fig.colorbar(c, shrink=0.85)
    cbar.ax.set_ylabel('Temperature (K)', rotation=90, fontsize=16)
    cbar.ax.tick_params(length=0, labelsize=16) 
    qv1=ax.quiver(xe[::6], ye[::6], ue[::6, ::6], ve[::6, ::6], pivot='middle', headwidth=2)
    ax.quiverkey(qv1, 0.8, 0.85, 10, r'10m/s', labelcolor=(0.3, 0.1, .2, 1),
                   labelpos='N', coordinates = 'figure', fontproperties={'size': 14, 'weight': 'bold'})
    plt.title('T$_{surf}$, winds at 1.3e4 Pa', fontsize=16)
    if tl_coord==True:
        plt.ylabel('Tidally-locked Latitude $\phi^{\prime} (^{\circ})$', fontsize=16)
        plt.xlabel('Tidally-locked Longitude $\lambda^{\prime} (^{\circ})$', fontsize=16)
    else:
        plt.ylabel('Latitude $\phi (^{\circ})$', fontsize=16)
        plt.xlabel('Longitude $\lambda (^{\circ})$', fontsize=16)
    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)
    plt.xlim(0,360)
    plt.ylim(-89,89)
    plt.show()

plot_tsurf(pcb_2d)
plot_tsurf(pcb_2d_regrid)