# Article title: Quantifying the Sampling Error on Burn Counts in Monte Carlo Wildfire Simulations using Poisson and Gamma Distributions.
# Journal name: Stochastic Environmental Research and Risk Assessment.
# Author names: Valentin Waeselynck, Gary Johnson, David Schmidt, Max Moritz, David Saah.
# Corresponding author: Valentin Waeselynck (vwaeselynck@sig-gis.com), affiliated to Spatial Informatics Group.

import numpy
import matplotlib.pyplot as pyplot
import scipy.stats

# GDAL commands used for: 
# 1) Before simulations, computing raster of prior parameters (prior mean at 10%/yr):
# gdal_calc.py --outfile pi0_yr-1.tiff --overwrite --type Float32 -A sierraville-gridfire-deck/MC_sampling_FM40_26910_clip.tif --calc '0.1*(1 - (A >= 90)*(A < 100))'
# 2) After simulations: extracting raster of burn counts burn_count.tiff (4936 simulated fires):
# gdal_calc.py --outfile burn_count.tiff --overwrite --type Float32 -A sierraville-gridfire-deck/outputs/burn_probability.tif --calc 'A * 4936.0'
# 3) Computing raster of Gamma posterior moments (20000 simulated years):
# gdal_calc.py --outfile gamma_posterior_burnfreq_yr-1.tiff --overwrite --type Float32 -A burn_count.tiff -B pi0_yr-1.tiff --calc '(1 - 1/(1 + 20000*B))*(1 + A)/20000' --calc '(1 - 1/(1 + 20000*B))*((1 + A)**0.5)/20000'

def plot_binomial_convergence():
    """
    Plots a histogram of Binomial samples along with the approximating Poisson PMF, at variance sample sizes.
    """
    n_samples = 30
    pi_C = 0.012345
    rho_C = 0.1
    sys = [350, 1000, 2500]
    fig, axs = pyplot.subplots(len(sys), 1, sharex=True, squeeze=False, figsize=(7, 2*len(sys)))
    #fig.subplots_adjust(hspace=0)
    x_max = 3.0*pi_C
    for sim_yrs, (ax,) in zip(sys, axs):
        ymax = 0.25
        ax.set_xlim([0, x_max])
        ax.set_ylim([0,ymax])
        ax.vlines(pi_C, 0, ymax, colors=['#2ca02c'], label=f"$\pi_C = $ {pi_C:.6f} yr$^{'{-1}'}$", ls="--")
        ax.set_xlabel("Burn frequency (yr$^{-1}$)")
        ax.set_ylabel("Prob. mass")

        nu_C = pi_C * sim_yrs
        N_I = int(nu_C / rho_C)

        n_poisson = scipy.stats.poisson(nu_C)
        pmf_ns = numpy.arange(0, x_max*sim_yrs, 1)
        n_pmf = n_poisson.pmf(pmf_ns)

        ax.plot(pmf_ns/sim_yrs, n_pmf, label="Poisson-based PMF", ls='-', marker='.', color='#1f77b4', linewidth=0.5)

        n_Cs = scipy.stats.binom.rvs(N_I, rho_C, size=n_samples)
        p_Cs = n_Cs / sim_yrs
        p_Cs_uniq, p_Cs_counts = numpy.unique(p_Cs, return_counts=True)

        XY = numpy.array([[p_C, (k-1)*0.08*ymax] for (p_C, n) in zip(p_Cs_uniq, p_Cs_counts) for k in range(1, n+1)])
        X = XY[:,0]
        Y = XY[:,1]

        ax.scatter(X,Y, label="$\hat{p}_C$ draws", cmap="Greys", facecolors='none', edgecolors='#ff7f0e')

        gamma_C = nu_C**-0.5
        ax.set_title(f"$N_I =$ {N_I} fires, $N_I/F_A =$ {sim_yrs:.0f} simulated years ($\gamma_C = $ {100*gamma_C:.0f}%)")
    #fig.suptitle("Illustration of the random variability of burn-frequency estimates in Monte-Carlo simulations.")
    fig.tight_layout()
    pyplot.legend()
    pyplot.show()


def plot_cv_from_simulated_years():
    fig, ax = pyplot.subplots(1, 1, squeeze=True, figsize=(6,3.5))
    sim_years = 10**numpy.arange(1, 5, 0.1)
    for piC in [1e-2, 5e-2, 1e-1]:
        cv = (sim_years * piC)**-0.5
        ax.plot(sim_years, cv, label=f"Burn-frequency $\\pi_C =$ {piC:.2E} yr$^{'{-1}'}$")
        ax.set_xlabel("Number of simulated years $N_I/F_A$ (yr)")
        ax.set_ylabel("Coefficient of Variation $\\gamma_C$")
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.legend()
    return fig


def plot_posterior_at_example_points(sim_yrs, pi0Cs, nCs, point_labels):
    """
    Plots the density of Gamma Posteriors for the population burn frequency of the given points.
    """
    #fig, rows = pyplot.subplots(int(len(nCs)/2), 2, squeeze=False, sharex=True, figsize = (13, 2.5*len(nCs)/2))
    fig, rows = pyplot.subplots(len(nCs), 1, squeeze=False, sharex=True, figsize = (8, 2.2*len(nCs)))
    axs = [ax for row in rows for ax in row]
    for ax, pi0C, nC, point_label in zip(axs, pi0Cs, nCs, point_labels):
        ax.set_xlabel("Asymptotic burn frequency $\pi_C$ (yr$^{-1}$)")
        ax.set_title(f"Location {point_label}: prior mean $\\pi^0_C = $ {pi0C} yr$^{'{-1}'}$, burn count $\\{'hat{n}'}_C = {nC}$")
        xmax = 2.2e-2
        xs = numpy.arange(0, xmax, 1e-4)
        gamma_k = 1 + nC
        gamma_theta = (1 - 1/(1 + sim_yrs*pi0C))/sim_yrs
        if gamma_theta > 0:
            ax.set_ylabel("Prob. density (yr)")
            gamma_post = scipy.stats.gamma(a=gamma_k, scale=gamma_theta)
            ax.plot(xs, gamma_post.pdf(xs), label="Posterior Probability Density Function (PDF)")
            mC = gamma_k * gamma_theta
            sC = gamma_k**0.5 * gamma_theta
            p99 = gamma_post.isf(1 - 0.99)
            ax.axvline(mC, 0, 1, ls="--", c='green', label=f"Posterior mean $\\{'hat{m}'}_C =$ {mC:.2E} yr$^{'{-1}'}$")
            ax.axvline(p99, 0, 1, ls=":", c='red', label=f"99th percentile p99 = {p99:.2E} yr$^{'{-1}'}$")
            gamma_mode = nC * gamma_theta
            ax.errorbar(gamma_mode, gamma_post.pdf(gamma_mode)/2, xerr=sC, label=f"Standard deviation $\\{'hat{s}'}_C =$ {sC:.2E} yr$^{'{-1}'}$")            
        else:
            ax.set_ylabel("Probability mass")
            ax.set_ylim([0, 1.2])
            ax.scatter([0.0], [1.0], label="Posterior Probability Mass Function (PMF)")
            ax.vlines(0, 0, 1, ls="-")
        ax.legend()
    #fig.suptitle(f"Posterior distributions of asymptotic burn-frequency at 4 example locations ({sim_yrs} simulated years)")
    fig.tight_layout()

#plot_posterior_at_example_points(20e3, [0, 1e-1, 1e-1, 1e-1], [0, 0, 158, 323], ["$C_0$", "$C_1$", "$C_2$", "$C_3$"])