import pandas as pd
from tqdm.auto import tqdm
import os
import glob
import numpy as np
from multiprocessing import Pool    
import polars as pl
import importlib
import re

import gc

import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib import colors
import matplotlib as mpl

import statsmodels.api as sm
import statsmodels.graphics.gofplots as gofplots

import pymc as pm
import arviz as az

import geopandas as gpd
import folium
from folium.features import GeoJsonTooltip
from shapely.geometry import Polygon, LineString, Point

%config InlineBackend.figure_format = 'retina'

# disable future warnings
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

sns.set_style("darkgrid")

pl.enable_string_cache()

!date

# the first year of data to use
start_year = 2010 
# the list of counties to use - FIPS codes
counties = [6085] 

# The minimum number of days of data of valid data for a station not to be dropped
MIN_DAYS_OF_DATA = 2 * 365

# Sampling parameters
# Whether to compute convergence checks
compute_convergence_checks = True
# Whether to compute the posterior predictive
sample_posterior_predictive = True
# Whether to compute the log likelihood
compute_log_likelihood = True
# The sampler to use
nuts_sampler='pymc'

# something to identify the run
nonce = "na"

print("Start year: ", start_year)
print("Counties: ", counties)

print("Convergence checks: ", compute_convergence_checks)
print("Sample posterior predictive: ", sample_posterior_predictive)
print("Compute log likelihood: ", compute_log_likelihood)
print("NUTS sampler: ", nuts_sampler)

print("Nonce: ", nonce)

!cd ../crates/py_analysis ; maturin develop -q --release

import py_analysis as pa
importlib.reload(pa.load)
importlib.reload(pa.display)
importlib.reload(pa.ghcn_d)
importlib.reload(pa.utils)
importlib.reload(pa)

geojson = pa.load_county_parquet('../data/georef-united-states-of-america-county.parquet')
geojson.head()

pa.ghcn_d.download_all_until(2022, directory="../data/ghcn_d/")
_ = pa.ghcn_d.ghcn_d_csv_to_parquet(verbose=True, files="../data/ghcn_d/*.csv.gz")

all_df = pl.scan_parquet('../data/ghcn_d/*.parquet').\
    filter(pl.col('date').dt.year() >= start_year)

all_df.fetch()

stations_df = pa.ghcn_d.load_stations(file="../data/stations_data.csv")
stations_df

# stations with data in the counties
q = stations_df.lazy()\
    .filter(pl.col('county').is_in(counties))\
    .join(all_df.filter(pl.col('element') == 'TMAX')\
                .group_by(pl.col('station')).agg(
                    [pl.col('station').count().alias('count'),
                     pl.col('q-flag').is_not_null().count().alias('null_q'),
                     pl.col('date').min().alias('min_date'),
                     pl.col('date').max().alias('max_date')]), 
          on='station', 
          how='left')\
    .filter(pl.col('count').is_not_null())\
    .sort('count', descending=True)\
    .lazy()

# q.show_graph()
counties_stations_df = q.collect()
counties_stations_df

def markers_fn(us_map): 
    return stations_df.filter(pl.col('county').is_in(counties)).to_pandas().apply(
        lambda row:folium.CircleMarker(location=[row['latitude'], row['longitude']], 
                                       radius=1, 
                                       hue='red', 
                                       fill=True).add_to(us_map), 
        axis=1)
pa.plot_county_geojson(geojson, counties, zoom_start=8, markers_fn=markers_fn)

# stations with data in the counties and no q-flag
q = stations_df.lazy() \
    .filter(pl.col('county').is_in(counties)) \
    .join(all_df.filter(pl.col('element') == 'TMAX') \
                .filter(pl.col('q-flag').is_null()) \
                .group_by(pl.col('station')).agg(pl.col('date')),
          on='station', 
          how='left') \
    .filter(pl.col('date').is_not_null())\
    .group_by(pl.col('station')).agg(pl.col('date').explode().alias('dates'))

counties_stations_activity_df = q.collect()
counties_stations_activity_df

def format_station_name(station_id, stations_df=stations_df, new_line=True):
    """
    Format the name of a station given its id.
    """
    station = stations_df.row(by_predicate=(pl.col('station') == station_id), named=True)
    if new_line:
        return f"{station['name']}\n{station_id}"
    else:
        return f"{station['name']} - {station_id}"

# heatmap with a row for each station and a column for each day
counties_stations_heatmap = counties_stations_activity_df.with_columns(pl.lit(1).alias('something'))\
    .explode('dates')\
    .rename({'dates': 'date'})\
    .pivot(index='date', columns=['station'], values='something')\
    .sort('date')\
    .fill_null(0)\
    .upsample(time_column='date', every='1d')\
    .sort('date')

# plot the heatmap with seaborn
fig, ax = plt.subplots(figsize=(20, 10))
sns.heatmap(counties_stations_heatmap.to_pandas().set_index('date').T, 
            ax=ax, cbar=False, cmap='Blues', xticklabels="auto", yticklabels="auto")
ax.set_xlabel('Date')
ax.set_ylabel('Station')
ax.set_title('County Stations Activity')
# only show the year on the x-axis

# ax.xaxis.set_major_locator(mdates.YearLocator())
ax.set_xticklabels([item.get_text()[:4] for item in ax.xaxis.get_ticklabels()])
# ylabel with station id and name
ax.set_yticklabels([format_station_name(item.get_text()) for item in ax.yaxis.get_ticklabels()])

del counties_stations_activity_df

plt.show()

stations_to_ignore = counties_stations_df.filter(pl.col('null_q') < MIN_DAYS_OF_DATA).select('station', 'null_q', 'min_date', 'max_date')
stations_to_ignore

q = stations_df.lazy()\
    .filter(pl.col('county').is_in(counties))\
    .filter(pl.col('station').is_in(list(stations_to_ignore['station'])).not_())\
    .join(all_df.filter(pl.col('element').is_in(['TMAX', 'TMIN']) & pl.col("q-flag").is_null())\
                .select([pl.col('station'), pl.col('element'), pl.col('date'), pl.col('value')]),
          on='station', 
          how='inner')\
    .with_columns([ 
        pl.col('station').cast(pl.Categorical),
        pl.col('state').cast(pl.Categorical),
        pl.col('name').cast(pl.Categorical),        
        pl.col('date').cast(pl.Date)
    ]) \
    .sort(['date', 'station', 'element'], descending=True) \
    .lazy()

wx_df = q.collect().pivot(index=['date', 'station', 'state', 'name', 'latitude', 'longitude', 'elevation', 'county'], columns='element', values='value')
print(f"wx_df.shape: {wx_df.shape}")
wx_df.head()

# dataframe with the data as index and a column for each station - each cell is the value of the station for that day
wx_df_pivot = wx_df.to_pandas().pivot(index='date', columns='station', values='TMAX')
# fill missing dates
wx_df_pivot = wx_df_pivot.asfreq('D', fill_value=np.nan)
# convert to int
wx_df_pivot = wx_df_pivot.astype(float)

# plot the heatmap with seaborn
fig, ax = plt.subplots(figsize=(20, 10))
sns.heatmap(wx_df_pivot.T, 
            ax=ax, cbar=True, 
            cmap='coolwarm', xticklabels="auto", yticklabels="auto")
ax.set_xlabel('Station')
ax.set_ylabel('Date')
ax.set_title('Stations Activity')
# background color to see missing data
ax.set_facecolor('lightyellow')

ax.set_xticklabels([item.get_text()[:4] for item in ax.xaxis.get_ticklabels()])
# yticklabel with station name
ax.set_yticklabels([format_station_name(item.get_text()) for item in ax.yaxis.get_ticklabels()])
# reduce the size of the yticklabels
plt.yticks(fontsize=8)


plt.title('Stations TMAX')
del wx_df_pivot
plt.show()

plt.figure(figsize=(20, 10))
sns.scatterplot(data=wx_df.to_pandas(), x='date', y='TMAX', hue='station', alpha=0.5, s=10)
# legend on two columns
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0., ncol=2)
plt.title('Stations TMAX')

# keep the color map 
handles, labels = plt.gca().get_legend_handles_labels()
station_cmap = {}
for h, l in zip(handles, labels):
    station_cmap[l] = h.get_color()

plt.show()

# grid of plots with a plot for each station
g = sns.FacetGrid(wx_df.to_pandas(), row="station", hue="station", aspect=5, height=1, palette=station_cmap, sharex=True)
g.map(sns.lineplot, "date", "TMAX", alpha=0.9, linewidth=0.5, zorder=100)
g.set_titles("{row_name}")
g.add_legend()

# g.fig.subplots_adjust(hspace=0.5)
g.fig.set_size_inches(10, 20)
plt.show()

wx_df_plus = wx_df.with_columns([
    pl.col('date').dt.year().alias('year'), 
    pl.col('date').dt.month().alias('month'),
    pl.col('date').dt.ordinal_day().alias('day_of_year')])
print(f"wx_df_plus.shape: {wx_df_plus.shape}")
del wx_df
wx_df_plus.head()

wx_pd = wx_df_plus.to_pandas()

plt.figure(figsize=(20, 10))
# add year and day_of_year columns
sns.scatterplot(data=wx_pd.reset_index(), x='day_of_year', y='TMAX', hue='year', style="station", 
                alpha=0.5, palette='RdBu_r', s=20)
# legend on two columns
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0., ncol=2)
ax = plt.gca()
ax.xaxis.set_major_locator(mdates.MonthLocator())
ax.xaxis.set_major_formatter(mdates.DateFormatter('%b'))
plt.title('Stations TMAX')
plt.show()

wx_df_plus_ = wx_df_plus.to_pandas()
wx_df_plus_ = wx_df_plus_.groupby([wx_df_plus_.date.dt.to_period('M'), 'station']).agg({'TMAX': ['mean', 'count']}).reset_index()
# filter out (stations, year) with less 20 points - per month
wx_df_plus_ = wx_df_plus_[wx_df_plus_[('TMAX', 'count')] > 20]
# remove multiindex
wx_df_plus_.columns = [' '.join(col).strip() for col in wx_df_plus_.columns.values]
# rename columns
wx_df_plus_.rename(columns={'TMAX count': 'count', 'TMAX mean': 'TMAX'}, inplace=True)
# drop year column
wx_df_plus_.drop(columns=['count'], inplace=True)
wx_df_plus_['date'] = wx_df_plus_['date'].astype('datetime64[ns]')
wx_df_plus_.head()

fig = plt.figure(figsize=(20, 10))
# add year and day_of_year columns
sns.scatterplot(wx_df_plus_, x='date', y='TMAX', hue='station', 
                alpha=0.5, palette=station_cmap, s=50)
# legend on two columns
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0., ncol=2)
# format the x-axis - automatic
fig.autofmt_xdate()

plt.title('Stations TMAX by Month')
plt.show()

del wx_df_plus_
del wx_df_plus

# remove rows with missing data for TMAX
wx_pd.dropna(subset=['TMAX'], inplace=True)

wx_pd.sort_values(['date', 'station'], inplace=True)
wx_pd.set_index(['date', 'station'], inplace=True)
print(f"wx_pd.shape: {wx_pd.shape}")
wx_pd.head()

station_idxs, stations = pd.factorize(wx_pd.index.get_level_values(1))
print(f"Number of stations: {len(stations)}")
print(f"Number of station_idxs: {len(station_idxs)}")
stations

time_idxs, times = pd.factorize(wx_pd.index.get_level_values(0))
print(f"Number of times: {len(times)}")
print(f"Number of time_idxs: {len(time_idxs)}")
times

times_f = (times.astype('datetime64[ns]').astype(np.int64) // 10**9).astype(np.float64) / 60 / 60 / 24 / 365.25
print(f"times_f.shape: {times_f.shape}")
times_f

import concurrent.futures

def sample(model, draws=500, tune=1000, chains=4, cores=4, 
           compute_convergence_checks=True,
           sample_posterior_predictive=True, 
           compute_log_likelihood=True, 
           nuts_sampler='pymc'):
    """
        Sample from a PyMC3 model and return the InferenceData and LOO.
    
        Parameters
        ----------
        model: PyMC3 model
        draws: int
            Number of draws from the posterior distribution
        tune: int
            Number of tuning steps
        chains: int
            Number of chains
        cores: int
            Number of cores
        sample_prior_predictive: bool
            Weheter to sample prior predictive
        compute_convergence_checks: bool
            Whether to compute convergence checks
        sample_posterior_predictive: bool
            Whether to sample posterior predictive
        compute_log_likelihood: bool
            Whether to compute log likelihood
        nuts_sampler: str
            NUTS sampler to use. One of 'pymc', 'blackjax', 'numpyro', 'nutpie'

        Returns
        -------
        idata: InferenceData
            InferenceData object
    """
    assert nuts_sampler in ['pymc', 'blackjax', 'numpyro', 'nutpie']

    with model:        

        print(f"Sampling {draws} draws, {tune} tune, {chains} chains, {cores} cores")
        idata = pm.sample(draws=draws, tune=tune, chains=chains, cores=cores, 
                          compute_convergence_checks=compute_convergence_checks,
                          return_inferencedata=True, nuts_sampler=nuts_sampler,
                          idata_kwargs={'log_likelihood': False})

        if sample_posterior_predictive:
            # sample posterior predictive        
            print(f"Sampling posterior predictive")
            pm.sample_posterior_predictive(idata, extend_inferencedata=True, return_inferencedata=True)

        if compute_log_likelihood:
            # compute log likelihood
            print(f"Computing log likelihood")
            pm.compute_log_likelihood(idata)

    # loo = az.loo(idata)    
    return idata # , loo

def run_model(fn, *args, **kwargs):
    """
        Run a function in a separate process and return the result.
        This is useful to avoid memory leaks in PyMC.
    """
    with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
        future = executor.submit(fn, *args, **kwargs)        
        return future.result()

def plot_summary(idata, var_names, var_posteriors=[], var_ppcs=[], sample_posterior_predictive=True):

    if not sample_posterior_predictive:
        var_ppcs = []
        print("No posterior predictive samples")

    # 1 subplot for the trace - 1 for each var of interest
    n_rows = len(var_names) + max(len(var_posteriors), len(var_ppcs))
    fig, axs = plt.subplots(n_rows, 2, figsize=(20, 5 * n_rows))

    # plot the trace
    az.plot_trace(idata, var_names=var_names, kind='rank_vlines', axes=np.array(axs[:len(var_names),:]).reshape(len(var_names), 2))    
    # increase spacing between plots    
    # plt.subplots_adjust(hspace=0.8)

    # plot the posterior
    for i, var_of_interest in enumerate(var_posteriors):
        az.plot_posterior(idata, var_names=var_of_interest, ax=axs[len(var_names) + i, 0], kind='kde')
        axs[i + 1, 0].set_title(f"{var_of_interest} posterior")
    
    # plot the posterior predictive    
    for i, var_of_interest in enumerate(var_ppcs):
        az.plot_ppc(idata, var_names=var_of_interest, ax=axs[len(var_names) + i, 1], kind='kde')
        axs[i + 1, 1].set_title(f"{var_of_interest} posterior predictive")

def keep_vars(idata, threshold=65):
    """
        Keep only the vars with total dimension smaller than threshold
    """
    vars = []
    for k in idata.posterior.data_vars:    

        # ignore 'chain' and 'draw' dimensions and keep only the ones with a product of the remaining dimensions <= threshold
        s = idata.posterior[k].shape
        sn = idata.posterior[k].dims
        shape = 1
        for d in sn:
            if d not in ['chain', 'draw']:
                shape *= s[sn.index(d)]
        if shape <= threshold:
            vars.append(k)
    return vars

def save_idata(idata, filename):
    """
        Save the InferenceData to a file
    """
    idata.to_netcdf(filename)

def load_idata(filename):
    """
        Load the InferenceData from a file
    """
    return az.from_netcdf(filename)

def stratify(df, by, target_samples=10000):
    """
        Stratify a dataframe by a column and sample each group to have at most target_samples rows.
    """    
    def sample_(x, target_samples=10000):
        if len(x) <= target_samples:
            return x
        else:
            return x.sample(target_samples)

    return df.groupby(by, group_keys=False).apply(sample_, target_samples=target_samples)

def plot_ppc(idata, max_years=5, target_samples=10_000, bw_adjust=.2, hist=False):
    """
        Plot the posterior predictive check for the model

        We want to compare the density of the posterior predictive with the density of the observed data. 
        We do this for some slices of our models, namely 'stations' and 'years'. 

        We plot the posterior predictive as a KDE and the observed data as a histogram.
        If they are similar, then the model is not too bad.

        Parameters
        ----------
        idata: InferenceData
            InferenceData object
        max_years: int
            Maximum number of years to plot for each station. If there are more years, we select evenly spaced years.            
        target_samples: int
            Maximum number of samples to plot for each station and year
        bw_adjust: float
            Bandwidth adjustment for the KDE - if hist=True
        hist: bool - default False
            Whether to plot the histogram of the observed data
    """
    
    # prepare the data
    ddf = idata.posterior_predictive.to_dataframe().reset_index()
    ddf['station'] = ddf['obs_id'].apply(lambda x: x.split('_')[0])
    ddf['year'] = ddf['obs_id'].apply(lambda x: x.split('_')[1]).astype(int)
    ddf.drop(columns=['obs_id', 'chain', 'draw'], inplace=True)

    # stratify the data by station and year
    if target_samples is not None:
        ddf = stratify(ddf, by=['station', 'year'], target_samples=target_samples)

    n_stations = ddf['station'].nunique()
    
    years = ddf['year'].unique()
    if len(years) > max_years:
        # select evenly spaced years
        years = years[np.linspace(0, len(years) - 1, 5, dtype=int)]
    n_years = len(years)

    fig, ax = plt.subplots(n_stations, n_years, figsize=(1 * n_years, 1 * n_stations), sharex=True, sharey=True)

    palette = sns.color_palette("viridis", ddf['year'].nunique()) 

    palette = {
        year: palette[i] for i, year in enumerate(ddf['year'].unique())
    }

    wx_pd_ = wx_pd.reset_index()

    for i, s in enumerate(ddf['station'].unique()):

        # filter the dataframe for the station
        df = ddf[ddf['station'] == s]

        wx_pd_s = wx_pd_[wx_pd_['station'] == s]

        for j, y in enumerate(years):
            df_y = df[df['year'] == y]
            wx_pd_s_y = wx_pd_s[wx_pd_s['year'] == y]
            if hist:
                if len(df_y) > 0:
                    # plot the posterior predictive - kde
                    sns.kdeplot(data=df_y, x='temperature', hue='year', ax=ax[i, j], palette=palette, legend=False, linewidth=0.5,  bw_adjust=bw_adjust)
                    # no x-axis label
                    ax[i, j].set_xlabel(None)
                    # no y-axis label

                # plot the observed data for the station - kde                
                if len(wx_pd_s_y) > 0:
                    # sns.kdeplot(data=wx_pd_s, x='TMAX', ax=ax[i, j], hue='year', palette=palette, linestyle='--', label='observed')
                    # hist of the observed data for the station                    
                    sns.histplot(data=wx_pd_s_y, x='TMAX', ax=ax[i, j], color='black', alpha=0.2, legend=False, stat='density')
                    # no x-axis label
                    ax[i, j].set_xlabel(None)
                    # no y-axis label
                    ax[i, j].set_ylabel(None)
            else:
                if len(df_y) > 0:
                    # qqplot
                    pp_1 = sm.ProbPlot(wx_pd_s_y['TMAX'], fit=False)
                    pp_2 = sm.ProbPlot(df_y['temperature'], fit=False)

                    gofplots.qqplot_2samples(pp_1, pp_2, line='45', ax=ax[i, j], xlabel=None, ylabel=None)                                         
                    
                    # remove labels    
                    ax[i, j].set_ylabel(None)
                    ax[i, j].set_xlabel(None)

            # title only for the first row
            if i == 0:
                ax[i, j].set_title(f"{y}")
            
        # set title for the row
        station_name = wx_pd_s['name'].iloc[0]
        # limit to 13 characters
        station_name = station_name[:13]
        ax[i, 0].set_ylabel(f"{s}\n{station_name}")
        # reduce the size of the label
        ax[i, 0].yaxis.label.set_size(8)
        
    # set title 
    if hist:
        fig.suptitle('Posterior Predictive (KDE) and Observations (Histogram)')
    else:
        fig.suptitle('QQplot of Observed and Posterior Predictive Samples')

    plt.tight_layout()
    # reduce the space between plots - horizontal and vertical
    plt.subplots_adjust(wspace=0.05, hspace=0.05)


def trend_station_marker(us_map, idata_df):
    """
        Define folium marker for stations
    """
    cmap = plt.get_cmap('RdBu_r')
    vmin = idata_df['mean'].min()
    vmax = idata_df['mean'].max()
    # symetrize the colorbar
    vmax = max(abs(vmin), abs(vmax))
    vmin = -vmax
    norm = colors.Normalize(vmin=vmin, vmax=vmax)
    for idx, row in idata_df.iterrows():
        latitude = row['latitude']
        longitude = row['longitude']   
        name = row['name']     
        color = cmap(norm(row['mean']))
        color = colors.rgb2hex(color)        
        tip = f"<u>id: {row['station']}</u><br/><b>{name}</b><br/>" + \
            f"mean Tmax trend: {row['mean']*10:.4f} C/decade<br/>" + \
            f"<b>95% CI: [{row['q_0.025']*10:.4f}, {row['q_0.975']*10:.4f}] C/decade</b><br/>" + \
            f"Data from {row['min_date']} to {row['max_date']}."
        folium.CircleMarker(location=[latitude, longitude], radius=9, tooltip=tip, color=color, fillColor=color, fillOpacity=0.8).add_to(us_map)
    return us_map

wx_pd.head()

def build_0(stations, wx_pd, time_idxs, times_f):
    coords={
        "station":stations,
        "month": ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"],
        "obs_id":[f'{loc}_{time.year}_{time.month}_{time.day}' for time, loc in wx_pd.index.values]
    }

    with pm.Model(coords=coords) as model:    
        t_ = pm.ConstantData('t', times_f)
        
        month = pm.ConstantData('month', wx_pd['month'].values, dims=['obs_id'])
        temp = pm.ConstantData('temp', wx_pd['TMAX'].values, dims=['obs_id'])

        # priors
        average_temperature = pm.Normal("average_temperature", 20, 10)    
        month_offset = pm.ZeroSumNormal("month_offset", sigma=5, dims="month")
        trend = pm.Normal("trend", 0, 1)        

        expected_temperature = pm.Deterministic(
            "expected_temperature", 
            average_temperature + 
            (trend * t_[time_idxs]) + 
            month_offset[month - 1],
            dims=("obs_id")
        )

        sigma = pm.HalfNormal("sigma", 5)
        # likelihood
        pm.Normal("temperature", mu=expected_temperature, sigma=sigma, observed=temp, dims=("obs_id"))

    return model

pm.model_to_graphviz(build_0(stations, wx_pd, time_idxs, times_f))

def build_and_sample(stations, wx_pd, time_idxs, times_f, 
                     compute_convergence_checks=True, 
                     sample_posterior_predictive=True, 
                     compute_log_likelihood=True, nuts_sampler='pymc'):
    model = build_0(stations, wx_pd, time_idxs, times_f)
    idata = sample(model, 
                     compute_convergence_checks=compute_convergence_checks, 
                     sample_posterior_predictive=sample_posterior_predictive, 
                     compute_log_likelihood=compute_log_likelihood, 
                     nuts_sampler=nuts_sampler)
    return idata

idata_0 = run_model(build_and_sample, stations, wx_pd, time_idxs, times_f, 
                    compute_convergence_checks=compute_convergence_checks, 
                    sample_posterior_predictive=sample_posterior_predictive, 
                    compute_log_likelihood=compute_log_likelihood, 
                    nuts_sampler=nuts_sampler)

idata_0

plot_ppc(idata_0)
plt.show()

az.summary(idata_0, var_names=['average_temperature', 'trend', 'sigma', 'month_offset'], round_to=3, kind='stats')

plot_summary(idata_0, 
             var_names=['average_temperature', 'trend', 'sigma', 'month_offset'], 
             var_posteriors=['trend'], 
             var_ppcs=['temperature'], 
             sample_posterior_predictive=sample_posterior_predictive)

az.plot_forest(idata_0, var_names=keep_vars(idata_0), 
               combined=True, kind='forestplot', hdi_prob=0.95, figsize=(5, 10))
plt.title('Model 0')
plt.show()

idata_0.posterior['total_trend'] = idata_0.posterior['trend'] 

def plot_posterior_predictive_model(wx_pd, idata, stations_df, station_cmap, sample_posterior_predictive=True):
    wx_df_pd = wx_pd[['TMAX']]
    new_index = pd.MultiIndex.from_product(wx_df_pd.index.levels)
    wx_df_pd = wx_df_pd.reindex(new_index).reset_index()
 
    stations = wx_df_pd['station'].unique()

    # grid of plots with a plot for each station
    g = sns.FacetGrid(wx_df_pd, row="station", hue="station", aspect=9, height=3, palette=station_cmap, sharex=True)
    g.map(pa.lineplot_breaknans, "date", "TMAX", alpha=0.9, linewidth=0.5, zorder=100)
    g.set_titles("{row_name}")
    g.add_legend()

    trend = idata.posterior['total_trend'].mean(dim=['chain', 'draw']).values
    trend_ci = az.hdi(idata.posterior['trend'], hdi_prob=0.95)
    trend_ci_lower = trend_ci.sel(hdi='lower').trend
    trend_ci_higher = trend_ci.sel(hdi='higher').trend

    # annotate the trend for each station
    for ax, station in zip(g.axes.flat, stations):        
        ax.annotate(f"trend: {trend*10:.3f}C/decade - 95%CI: [{trend_ci_lower*10:.3f}, {trend_ci_higher*10:.3f}]C/decade",
                    xy=(0.05, 0.95), 
                    xycoords='axes fraction', fontsize=12,
                    horizontalalignment='left', verticalalignment='top')
        
        # plot the posterior predictive 
        if sample_posterior_predictive:
            obs_ids = [f'{loc}_{time.year}_{time.month}_{time.day}' for time, loc in wx_pd.index.values if loc == station]
            all_times = wx_pd.index.get_level_values(0).unique()
            times_station = [time for time, loc in wx_pd.index.values if loc == station]    
            mean_ppc = idata.posterior_predictive['temperature'].sel(obs_id=obs_ids).mean(dim=['chain', 'draw']).values
            # 95%-ci
            ci_ppc = az.hdi(idata.posterior_predictive['temperature'].sel(obs_id=obs_ids), hdi_prob=0.95)
            lower_ci = ci_ppc.sel(hdi='lower').temperature.values
            higher_ci = ci_ppc.sel(hdi='higher').temperature.values

            # build dataframe with mean_ppc, lower_ci and higher_ci
            station_trend_df = pd.DataFrame({'mean_ppc': mean_ppc, 'lower_ci': lower_ci, 'higher_ci': higher_ci}, index=times_station)
            # resample to all_times
            station_trend_df = station_trend_df.reindex(all_times)

            # plot the mean_ppc
            ax.plot(station_trend_df.index, station_trend_df['mean_ppc'], color='red', linewidth=0.5, linestyle='--')

            # plot the 95%-CI
            ax.fill_between(station_trend_df.index, station_trend_df['lower_ci'], station_trend_df['higher_ci'], color='black', alpha=0.2, linewidth=0.5)

        # set the title for each plot with the station id and name        
        ax.set_title(format_station_name(station, new_line=False))

    return g

g = plot_posterior_predictive_model(wx_pd, idata_0, stations_df, station_cmap, sample_posterior_predictive=sample_posterior_predictive)
plt.show()

def build_1(stations, wx_pd, time_idxs, times_f, station_idxs):
    coords={
    "station": stations,
    "month": ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"],
    "obs_id":[f'{loc}_{time.year}_{time.month}_{time.day}' for time, loc in wx_pd.index.values]
    }

    with pm.Model(coords=coords) as model:        
        t_ = pm.ConstantData('t', times_f)
        
        month = pm.ConstantData('month', wx_pd['month'].values, dims=['obs_id'])
        temp = pm.ConstantData('temp', wx_pd['TMAX'].values, dims=['obs_id'])

        # priors
        average_temperature = pm.Normal("average_temperature", 20, 10)    
        month_offset = pm.ZeroSumNormal("month_offset", sigma=5, dims="month")
        trend = pm.Normal("trend", 0, 1)        

        station_offset = pm.ZeroSumNormal("station_offset", sigma=3.0, dims="station")

        expected_temperature = pm.Deterministic(
            "expected_temperature", 
            average_temperature + 
            (trend * t_[time_idxs]) + 
            month_offset[month - 1] + 
            station_offset[station_idxs],
            dims=("obs_id")
        )

        sigma = pm.HalfNormal("sigma", 5)
        # likelihood
        pm.Normal("temperature", mu=expected_temperature, sigma=sigma, observed=temp, dims=("obs_id"))
    
    return model

pm.model_to_graphviz(build_1(stations, wx_pd, time_idxs, times_f, station_idxs))

def build_and_sample(stations, wx_pd, time_idxs, times_f, station_idxs, 
                     compute_convergence_checks=True, 
                     sample_posterior_predictive=True, 
                     compute_log_likelihood=True, nuts_sampler='pymc'):
    model = build_1(stations, wx_pd, time_idxs, times_f, station_idxs)
    idata = sample(model, 
                     compute_convergence_checks=compute_convergence_checks, 
                     sample_posterior_predictive=sample_posterior_predictive, 
                     compute_log_likelihood=compute_log_likelihood, 
                     nuts_sampler=nuts_sampler)
    return idata

idata_1 = run_model(build_and_sample, stations, wx_pd, time_idxs, times_f, station_idxs,
                    compute_convergence_checks=compute_convergence_checks, 
                    sample_posterior_predictive=sample_posterior_predictive, 
                    compute_log_likelihood=compute_log_likelihood, 
                    nuts_sampler=nuts_sampler)

idata_1

plot_ppc(idata_1)
plt.show()

az.summary(idata_1, var_names=['average_temperature', 'trend', 'sigma', 'month_offset', 'station_offset'], round_to=3, kind='stats')

plot_summary(idata_1, 
             var_names=['average_temperature', 'trend', 'sigma', 'month_offset', 'station_offset'], 
             var_posteriors=['trend'], 
             var_ppcs=['temperature'], 
             sample_posterior_predictive=sample_posterior_predictive)

az.plot_forest(idata_1, var_names=keep_vars(idata_1), 
               combined=True, kind='forestplot', hdi_prob=0.95, figsize=(10, 20))

plt.axvline(x=0, color='red', linestyle='--')

plt.title('Model 1')
plt.show()

idata_1.posterior['total_trend'] = idata_1.posterior['trend'] 

g = plot_posterior_predictive_model(wx_pd, idata_1, stations_df, station_cmap, sample_posterior_predictive=sample_posterior_predictive)
plt.show()

def build_2(stations, wx_pd, time_idxs, times_f, station_idxs):
    coords={
    "station": stations,
    "month": ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"],
    "obs_id":[f'{loc}_{time.year}_{time.month}_{time.day}' for time, loc in wx_pd.index.values]
    }

    with pm.Model(coords=coords) as model:
        t_ = pm.ConstantData('t', times_f)
        
        month = pm.ConstantData('month', wx_pd['month'].values, dims=['obs_id'])
        temp = pm.ConstantData('temp', wx_pd['TMAX'].values, dims=['obs_id'])

        # priors
        average_temperature = pm.Normal("average_temperature", 20, 10)    
        month_offset = pm.ZeroSumNormal("month_offset", sigma=5, dims="month")
    
        trend_mu = pm.Normal("trend_mu", 0, 1)
        trend_sigma = pm.HalfNormal("trend_sigma", 1)
        trend_month = pm.Normal("trend_month", mu=trend_mu, sigma=trend_sigma, dims="month")

        station_offset = pm.ZeroSumNormal("station_offset", sigma=3.0, dims="station")    
        
        expected_temperature = pm.Deterministic(
            "expected_station_temperature", 
            average_temperature + 
            (trend_month[month - 1] * t_[time_idxs]) + 
            month_offset[month - 1] + 
            station_offset[station_idxs],
            dims=("obs_id")
        )

        sigma = pm.HalfNormal("sigma", 5)
        # likelihood
        pm.Normal("temperature", mu=expected_temperature, sigma=sigma, observed=temp, dims=("obs_id"))
    
    return model

pm.model_to_graphviz(build_2(stations, wx_pd, time_idxs, times_f, station_idxs))

def build_and_sample(stations, wx_pd, time_idxs, times_f, station_idxs, 
                     compute_convergence_checks=True, 
                     sample_posterior_predictive=True, 
                     compute_log_likelihood=True, nuts_sampler='pymc'):
    model = build_2(stations, wx_pd, time_idxs, times_f, station_idxs)
    idata = sample(model, 
                     compute_convergence_checks=compute_convergence_checks, 
                     sample_posterior_predictive=sample_posterior_predictive, 
                     compute_log_likelihood=compute_log_likelihood, 
                     nuts_sampler=nuts_sampler)
    return idata

idata_2 = run_model(build_and_sample, stations, wx_pd, time_idxs, times_f, station_idxs,
                    compute_convergence_checks=compute_convergence_checks, 
                    sample_posterior_predictive=sample_posterior_predictive, 
                    compute_log_likelihood=compute_log_likelihood, 
                    nuts_sampler=nuts_sampler)

idata_2

plot_ppc(idata_2)
plt.show()

az.summary(idata_2, 
           var_names=['average_temperature', 'trend_mu', 'trend_month', 
                      'trend_sigma', 'sigma', 'month_offset', 'station_offset'], 
           round_to=3, kind='stats')

plot_summary(idata_2, 
             var_names=['average_temperature', 'trend_mu', 'trend_month', 
                        'trend_sigma', 'month_offset', 'station_offset'],              
             var_posteriors=['trend_mu', 'trend_month'],
             var_ppcs=['temperature'],
            sample_posterior_predictive=sample_posterior_predictive)

az.plot_forest(idata_2, var_names=keep_vars(idata_2), 
               combined=True, kind='forestplot', hdi_prob=0.95, figsize=(10, 20))
# vertical dash red line at 0
plt.axvline(x=0, color='red', linestyle='--')

plt.title('Model 2')
plt.show()

idata_2.posterior['total_trend'] = idata_2.posterior['trend_month'] 

def plot_posterior_predictive_model_alt(wx_pd, idata, stations_df, station_cmap, sample_posterior_predictive=True, color_trend=True):
    wx_df_pd = wx_pd[['TMAX']]
    new_index = pd.MultiIndex.from_product(wx_df_pd.index.levels)
    wx_df_pd = wx_df_pd.reindex(new_index).reset_index()

    stations = wx_df_pd['station'].unique()

    # grid of plots with a plot for each station
    g = sns.FacetGrid(wx_df_pd, row="station", hue="station", aspect=9, height=3, palette=station_cmap, sharex=True)
    g.map(pa.lineplot_breaknans, "date", "TMAX", alpha=0.9, linewidth=0.5, zorder=100)
    g.set_titles("{row_name}")
    g.add_legend()

    cmap = plt.get_cmap('RdBu_r')
    trend_min = idata.posterior['total_trend'].min().values
    trend_max = idata.posterior['total_trend'].max().values
    # symetrize the norm
    trend_max = max(abs(trend_min), abs(trend_max))
    trend_min = -trend_max
    norm = colors.Normalize(vmin=trend_min, vmax=trend_max)

    coords = list(idata.posterior['total_trend'].coords)

    if 'month' in coords:
        trend = idata.posterior['total_trend'].mean(dim=['chain', 'draw']).mean(dim=['month'])
        trend_ci_lower = idata.posterior['total_trend'].quantile(0.025, dim=['chain', 'draw']).mean(dim=['month'])
        trend_ci_higher = idata.posterior['total_trend'].quantile(0.975, dim=['chain', 'draw']).mean(dim=['month'])
    else:
        trend = idata.posterior['total_trend'].mean(dim=['chain', 'draw'])
        trend_ci_lower = idata.posterior['total_trend'].quantile(0.025, dim=['chain', 'draw'])
        trend_ci_higher = idata.posterior['total_trend'].quantile(0.975, dim=['chain', 'draw'])    

    # annotate the trend for each station
    for ax, station in zip(g.axes.flat, stations):
        if 'station' in coords:
            trend_ = trend.sel(station=station).values
            trend_ci_lower_ = trend_ci_lower.sel(station=station).values
            trend_ci_higher_ = trend_ci_higher.sel(station=station).values
        else:
            trend_ = trend.values
            trend_ci_lower_ = trend_ci_lower.values
            trend_ci_higher_ = trend_ci_higher.values
        ax.annotate(f"trend (averaged over months): {trend_*10:.3f}C/decade - 95%CI: [{trend_ci_lower_*10:.3f}, {trend_ci_higher_*10:.3f}]C/decade",
                    xy=(0.05, 0.95), xycoords='axes fraction', fontsize=12,
                    horizontalalignment='left', verticalalignment='top')
        
        # plot the posterior predictive 
        if sample_posterior_predictive:
            obs_ids = [f'{loc}_{time.year}_{time.month}_{time.day}' for time, loc in wx_pd.index.values if loc == station]
            all_times = wx_pd.index.get_level_values(0).unique()
            times_station = [time for time, loc in wx_pd.index.values if loc == station]    
            mean_ppc = idata.posterior_predictive['temperature'].sel(obs_id=obs_ids).mean(dim=['chain', 'draw']).values
            # 95%-ci
            ci_ppc = az.hdi(idata.posterior_predictive['temperature'].sel(obs_id=obs_ids), hdi_prob=0.95)
            lower_ci = ci_ppc.sel(hdi='lower').temperature.values
            higher_ci = ci_ppc.sel(hdi='higher').temperature.values

            # build dataframe with mean_ppc, lower_ci and higher_ci
            station_trend_df = pd.DataFrame({'mean_ppc': mean_ppc, 'lower_ci': lower_ci, 'higher_ci': higher_ci}, index=times_station)
            # resample to all_times
            station_trend_df = station_trend_df.reindex(all_times)

            # plot the mean_ppc
            ax.plot(station_trend_df.index, station_trend_df['mean_ppc'], color='red', linewidth=0.5, linestyle='--')

            # plot the 95%-CI
            ax.fill_between(station_trend_df.index, station_trend_df['lower_ci'], station_trend_df['higher_ci'], color='black', alpha=0.2, linewidth=0.5)

        if color_trend:
            # change the background color of each plot with the trend        
            ax.set_facecolor(cmap(norm(trend_)))

        # set the title for each plot with the station id and name    
        ax.set_title(format_station_name(station, new_line=False))

    return g

# don't color the background with the trend since they are the same for all stations with this model.
g = plot_posterior_predictive_model_alt(wx_pd, idata_2, stations_df, station_cmap, sample_posterior_predictive=sample_posterior_predictive, color_trend=False)
plt.show()

def build_3(stations, wx_pd, time_idxs, times_f, station_idxs):
    coords={
    "station": stations,
    "month": ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"],
    "obs_id":[f'{loc}_{time.year}_{time.month}_{time.day}' for time, loc in wx_pd.index.values]
    }

    with pm.Model(coords=coords) as model:
        t_ = pm.ConstantData('t', times_f)
        
        month = pm.ConstantData('month', wx_pd['month'].values, dims=['obs_id'])
        temp = pm.ConstantData('temp', wx_pd['TMAX'].values, dims=['obs_id'])

        # priors
        average_temperature = pm.Normal("average_temperature", 20, 10)    
        month_offset = pm.ZeroSumNormal("month_offset", sigma=5, dims="month")
    
        trend_mu = pm.Normal("trend_mu", 0, 1)
        trend_sigma = pm.HalfNormal("trend_sigma", 1)
        trend_station = pm.Normal("trend_station", mu=trend_mu, sigma=trend_sigma, dims="station")

        station_offset = pm.ZeroSumNormal("station_offset", sigma=3.0, dims="station")    
        
        expected_temperature = pm.Deterministic(
            "expected_temperature", 
            average_temperature +             
            month_offset[month - 1] + 
            (trend_station[station_idxs] * t_[time_idxs]) + 
            station_offset[station_idxs],
            dims=("obs_id")
        )

        sigma = pm.HalfNormal("sigma", 5)
        # likelihood
        pm.Normal("temperature", mu=expected_temperature, sigma=sigma, observed=temp, dims=("obs_id"))
    
    return model

pm.model_to_graphviz(build_3(stations, wx_pd, time_idxs, times_f, station_idxs))

def build_and_sample(stations, wx_pd, time_idxs, times_f, station_idxs, 
                     compute_convergence_checks=True, 
                     sample_posterior_predictive=True, 
                     compute_log_likelihood=True, nuts_sampler='pymc'):
    model = build_3(stations, wx_pd, time_idxs, times_f, station_idxs)
    idata = sample(model, 
                     compute_convergence_checks=compute_convergence_checks, 
                     sample_posterior_predictive=sample_posterior_predictive, 
                     compute_log_likelihood=compute_log_likelihood, 
                     nuts_sampler=nuts_sampler)
    return idata

idata_3 = run_model(build_and_sample, stations, wx_pd, time_idxs, times_f, station_idxs,
                    compute_convergence_checks=compute_convergence_checks, 
                    sample_posterior_predictive=sample_posterior_predictive, 
                    compute_log_likelihood=compute_log_likelihood, 
                    nuts_sampler=nuts_sampler)

idata_3

plot_ppc(idata_3)
plt.show()

az.summary(idata_3, 
           var_names=['average_temperature', 'trend_mu', 'trend_station', 
                      'sigma', 'month_offset', 'station_offset'], round_to=3, kind='stats')

plot_summary(idata_3, 
             var_names=['average_temperature', 'trend_mu', 'trend_station', 'sigma', 'month_offset', 'station_offset'], 
             var_posteriors=['trend_mu', 'trend_station'],
             var_ppcs=['temperature'],
            sample_posterior_predictive=sample_posterior_predictive)

az.plot_forest(idata_3, var_names=keep_vars(idata_3), 
               combined=True, kind='forestplot', hdi_prob=0.95, figsize=(10, 20))

plt.axvline(x=0, color='red', linestyle='--')

plt.title('Model 3')
plt.show()

idata_3.posterior['total_trend'] = idata_3.posterior['trend_station'] 

def plot_forest_trend_station(idata, counties_stations_heatmap, stations_df, title):
    # left subplot with the forest for trend_station, right subplot with the activity heatmap for the stations (same order)
    fig, axs = plt.subplots(1, 2, figsize=(20, 10))
    az.plot_forest(idata, var_names=['total_trend'], combined=True, kind='forestplot', hdi_prob=0.95, figsize=(20, 10), ax=axs[0])
    # red dotted vertical line for the 0 value
    axs[0].axvline(0, color='red', linestyle='--')

    axs[0].set_title(title)

    # plot the heatmap with seaborn

    # get the order of the stations in the forest plot - got to extract the station id from the label with a regex (between [ and ])
    station_order = [re.search(r'\[(\w+)\]', item.get_text()).group(1) for item in axs[0].get_yticklabels()]
    # reversed
    station_order = station_order[::-1]

    counties_stations_heatmap_ = counties_stations_heatmap.to_pandas()
    counties_stations_heatmap_.set_index('date', inplace=True)
    # reorder the columns
    counties_stations_heatmap_ = counties_stations_heatmap_[station_order]

    sns.heatmap(counties_stations_heatmap_.T, 
                ax=axs[1], cbar=False, cmap='Blues', xticklabels="auto", yticklabels="auto")
    axs[1].set_xlabel('Date')
    axs[1].set_ylabel('Station')
    axs[1].set_title('Stations Activity')
    # only show the year on the x-axis
    axs[1].set_xticklabels([item.get_text()[:5] for item in axs[1].xaxis.get_ticklabels()])
    # set the y-axis labels with the station id and name    
    axs[1].set_yticklabels([f"{station}\n{stations_df.row(by_predicate=(pl.col('station') == station), named=True)['name']}" for station in station_order])
    axs[1].tick_params(axis='y', labelsize=8)
    plt.show()

plot_forest_trend_station(idata_3, counties_stations_heatmap, stations_df, title='Model 3')

g = plot_posterior_predictive_model_alt(wx_pd, idata_3, stations_df, station_cmap, sample_posterior_predictive=sample_posterior_predictive)
plt.show()

def summarize_model(idata, counties_stations_df, model_name):
    # 95%-ci
    idata_df_q = idata.posterior['total_trend'].quantile([0.025, 0.975], dim=['chain', 'draw']).to_dataframe().reset_index().pivot(index='station', columns='quantile', values='total_trend').rename(columns={0.025: 'q_0.025', 0.975: 'q_0.975'})
    idata_df_m = idata.posterior['total_trend'].mean(dim=['chain', 'draw']).to_dataframe().reset_index().rename(columns={'total_trend': 'mean'}).set_index('station')
    idata_df = pd.concat([idata_df_q, idata_df_m], axis=1)
    counties_stations_df_pd = counties_stations_df.to_pandas()
    idata_df = idata_df.merge(counties_stations_df_pd[['station', 'latitude', 'longitude', 'name', 'min_date', 'max_date']], left_index=True, right_on='station')

    idata_df['month'] = np.nan
    idata_df['model'] = model_name

    del idata_df_q
    del idata_df_m

    return idata_df

idata3_df = summarize_model(idata_3, counties_stations_df, model_name='model_3')
idata3_df

# plot map with the mean trend_station
us_map = pa.plot_county_geojson(geojson, counties,
                                zoom_start=8,
                                markers_fn=lambda m: trend_station_marker(m, idata_df=idata3_df))
us_map

def build_4(stations, wx_pd, time_idxs, times_f, station_idxs):
    coords={
    "station": stations,
    "month": ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"],
    "obs_id":[f'{loc}_{time.year}_{time.month}_{time.day}' for time, loc in wx_pd.index.values]
    }

    with pm.Model(coords=coords) as model:
        t_ = pm.ConstantData('t', times_f)
        
        month = pm.ConstantData('month', wx_pd['month'].values, dims=['obs_id'])
        temp = pm.ConstantData('temp', wx_pd['TMAX'].values, dims=['obs_id'])

        # priors
        average_temperature = pm.Normal("average_temperature", 20, 10)    
        month_offset = pm.ZeroSumNormal("month_offset", sigma=5, dims="month")
    
        trend = pm.Normal("trend", 0, 1)

        trend_offset_station_sigma = pm.HalfNormal("trend_offset_station_sigma", 1)
        trend_offset_month_sigma = pm.HalfNormal("trend_offset_month_sigma", 1)
        trend_offset_station = pm.Normal("trend_offset_station", mu=0, sigma=trend_offset_station_sigma, dims=["station"])
        trend_offset_month = pm.Normal("trend_offset_month", mu=0, sigma=trend_offset_month_sigma, dims=["month"])

        station_offset = pm.ZeroSumNormal("station_offset", sigma=3.0, dims="station")    

        expected_temperature = pm.Deterministic(
            "expected_temperature", 
            average_temperature +             
            month_offset[month - 1] + 
            ((trend + trend_offset_station[station_idxs] + trend_offset_month[month - 1]) * t_[time_idxs]) + 
            station_offset[station_idxs],
            dims=("obs_id")
        )

        sigma = pm.HalfNormal("sigma", 5)
        # likelihood
        pm.Normal("temperature", mu=expected_temperature, sigma=sigma, observed=temp, dims=("obs_id"))
    
    return model    

pm.model_to_graphviz(build_4(stations, wx_pd, time_idxs, times_f, station_idxs))

def build_and_sample(stations, wx_pd, time_idxs, times_f, station_idxs, 
                     compute_convergence_checks=True, 
                     sample_posterior_predictive=True, 
                     compute_log_likelihood=True, nuts_sampler='pymc'):
    model = build_4(stations, wx_pd, time_idxs, times_f, station_idxs)
    idata = sample(model, 
                     compute_convergence_checks=compute_convergence_checks, 
                     sample_posterior_predictive=sample_posterior_predictive, 
                     compute_log_likelihood=compute_log_likelihood, 
                     nuts_sampler=nuts_sampler)
    return idata

idata_4 = run_model(build_and_sample, stations, wx_pd, time_idxs, times_f, station_idxs,
                    compute_convergence_checks=compute_convergence_checks, 
                    sample_posterior_predictive=sample_posterior_predictive, 
                    compute_log_likelihood=compute_log_likelihood, 
                    nuts_sampler=nuts_sampler)

idata_4

plot_ppc(idata_4)
plt.show()

az.summary(idata_4,
           var_names=['average_temperature', 'trend', 'trend_offset_station', 'trend_offset_month', 'sigma', 'month_offset', 'station_offset'], 
           round_to=3, kind='stats')

plot_summary(idata_4,
                var_names=['average_temperature', 'trend', 'trend_offset_station', 'trend_offset_month', 'sigma', 'month_offset', 'station_offset'], 
                var_posteriors=['trend', 'trend_offset_station', 'trend_offset_month'],
                var_ppcs=['temperature'],
                sample_posterior_predictive=sample_posterior_predictive)

az.plot_forest(idata_4, var_names=keep_vars(idata_4), 
               combined=True, kind='forestplot', hdi_prob=0.95, figsize=(5, 15))

idata_4.posterior['total_trend'] = idata_4.posterior['trend'] + idata_4.posterior['trend_offset_month'] + idata_4.posterior['trend_offset_station']

def plot_forest_trend_station_month(idata, counties_stations_heatmap, stations_df, title):
    # left subplot with the forest for trend_station, right subplot with the activity heatmap for the stations (same order)
    fig, axs = plt.subplots(1, 2, figsize=(16, 15))
    # compute the CI for the sum of the trend and trend_offset     
    months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
    idata_s = [idata.posterior['total_trend'].sel(month=m) for m in months]

    az.plot_forest(idata_s, var_names=['total_trend'],
                model_names=months,               
                combined=True, kind='forestplot', hdi_prob=0.95, figsize=(8, 15), ax=axs[0])

    axs[0].axvline(0, color='red', linestyle='--')

    axs[0].set_title(title)

    # plot the heatmap with seaborn
    # get the order of the stations in the forest plot - got to extract the station id from the label with a regex (between [ and ])
    station_order = [re.search(r'\[(\w+)\]', item.get_text()).group(1) for item in axs[0].get_yticklabels()]
    # reversed
    station_order = station_order[::-1]

    counties_stations_heatmap_ = counties_stations_heatmap.to_pandas()
    counties_stations_heatmap_.set_index('date', inplace=True)
    # reorder the columns
    counties_stations_heatmap_ = counties_stations_heatmap_[station_order]

    sns.heatmap(counties_stations_heatmap_.T, 
                ax=axs[1], cbar=False, cmap='Blues', xticklabels="auto", yticklabels="auto")
    axs[1].set_xlabel('Date')
    axs[1].set_ylabel('Station')
    axs[1].set_title('Stations Activity')
    # only show the year on the x-axis
    axs[1].set_xticklabels([item.get_text()[:5] for item in axs[1].xaxis.get_ticklabels()])
    # add station name to the y-axis
    axs[1].set_yticklabels([f"{x}\n{stations_df.row(by_predicate=(pl.col('station') == x), named=True)['name'][:13]}" for x in station_order])
    # reduce the font size
    axs[1].tick_params(axis='y', which='major', labelsize=6)

plot_forest_trend_station_month(idata_4, counties_stations_heatmap, stations_df, title="Model 4")
plt.show()

# dataframe with mean and CI for trend_sum
def summarize_model_with_month(idata, counties_stations_df, model_name):
    idata_df = idata.posterior['total_trend'].mean(dim=['chain', 'draw']).to_dataframe().reset_index().rename(columns={'total_trend': 'mean'})
    idata_df_ci = az.hdi(idata.posterior['total_trend'], hdi_prob=0.95).to_dataframe().reset_index()\
                    .pivot(index=['station', 'month'], columns='hdi', values='total_trend')\
                        .rename(columns={'lower': 'q_0.025', 'higher': 'q_0.975'})
    idata_df = idata_df.merge(idata_df_ci, left_on=['station', 'month'], right_index=True)    
    # merge with counties_stations_df
    idata_df = idata_df.merge(counties_stations_df.to_pandas()[['station', 'latitude', 'longitude', 'name', 'min_date', 'max_date']], left_on='station', right_on='station')
    idata_df['model'] = model_name

    return idata_df

idata4_df = summarize_model_with_month(idata_4, counties_stations_df, model_name='model_4')

# plot the posterior trend_offset 
fig, ax = plt.subplots(figsize=(10, 5))

# line plot with error bar with matplotlib and idata4_df
for station in idata4_df['station'].unique():
    color = station_cmap[station]
    # filter the dataframe for the station
    idata_4_total_trend_ = idata4_df[idata4_df['station'] == station]
    # plot the mean
    ax.plot(idata_4_total_trend_['month'], idata_4_total_trend_['mean'], color=color, linewidth=1, label=station)
    # plot the 95%-CI
    ax.fill_between(idata_4_total_trend_['month'], idata_4_total_trend_['q_0.025'], idata_4_total_trend_['q_0.975'], color=color, alpha=0.1, linewidth=0.5)
    
# y axis label
ax.set_ylabel('Total trend [C/year]')

plt.title('Model 4: Total trend (trend + trend_offset_station + trend_offset_month) - mean and 95%-CI)')

# legend on the right
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0., ncol=2)

# horizontal red line at 0
plt.axhline(y=0, color='red', linestyle='--')
plt.show()

months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
total_trend_df = idata_4.posterior['total_trend'].mean(dim=['chain', 'draw']).to_dataframe().reset_index().rename(columns={'total_trend': 'mean'}).pivot(index='station', columns='month', values='mean')
# sort columns by month
total_trend_df = total_trend_df[months]

# heatmap with seaborn
n_stations = len(total_trend_df.index)
fig, ax = plt.subplots(figsize=(10, n_stations))
sns.heatmap(total_trend_df, ax=ax, cmap='RdBu_r', center=0, annot=True, fmt='.3f', linewidths=.5)
# y axis tickslabel with station names
ax.set_yticklabels(total_trend_df.index.map(lambda x: f"{x}\n{stations_df.row(by_predicate=(pl.col('station') == x), named=True)['name']}"))

# colorbar legend
ax.collections[0].colorbar.set_label("Total trend [C/year]")

plt.title('Model 4: Total trend (trend + trend_offset_station + trend_offset_month) - mean')
plt.show()