

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt

def plot_ts(
    seas5, 
    era5,
    title: str = None
    ) -> None:

    fig = plt.figure(figsize=(10,5))

    seas5.plot(
        hue='ensemble', x='year', color="grey", label='', alpha=0.5)

    seas5.mean(dim="ensemble").plot(
        x='year', color="black", label='SEAS5', linewidth=3)

    xr.apply_ufunc(
        lambda x: ml4xp.tools.loess_ts(x, window=40, degree=2),
            seas5.mean(dim="ensemble"),
            input_core_dims = [["year"]],
            output_core_dims=[["year"]],
            vectorize=True,
            dask='parallelized',
            output_dtypes=[float]
            ).plot(
            x='year', color="black", label='', linewidth=3, linestyle='--')

    era5.plot(
        x='year', color="red", label='ERA5', linewidth=3)
    
    xr.apply_ufunc(
        lambda x: ml4xp.tools.loess_ts(x, window=40, degree=2),
            era5,
            input_core_dims = [["year"]],
            output_core_dims=[["year"]],
            vectorize=True,
            dask='parallelized',
            output_dtypes=[float]
            ).plot(
            x='year', color="red", label='', linewidth=3, linestyle='--')

    corr  = xr.corr(seas5.mean(dim="ensemble"), era5, dim="year").values
    try:
        corr = corr[0]
    except:
        pass
    
    # hide legend
    plt.legend()
    if title is not None:
        plt.title(title + f" | R = {corr:.2f}")

        
