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

def _bootstrap_plot(
    boot_vals: np.ndarray,
    pval: float,
    random_boot_vals: np.ndarray = None,
    bins: int = 50
    ) -> None:

    boot_vals.plot.hist(bins=bins, alpha=0.5, color="orange")

    plt.axvline(boot_vals.quantile(pval), color="red", linestyle="--")
    plt.axvline(boot_vals.quantile(1 - pval), color="red", linestyle="--")
    plt.axvline(boot_vals.median(), color="black", linestyle="--")
    
    plt.text(
        boot_vals.median(), 0, 
        str(boot_vals.median().values.round(2)),
        verticalalignment="bottom"
        )

    if random_boot_vals is not None:
        random_boot_vals.plot.hist(bins=bins, alpha=0.5, color="blue")
        plt.axvline(random_boot_vals.quantile(pval), color="blue", linestyle="--")
        plt.axvline(random_boot_vals.quantile(1 - pval), color="blue", linestyle="--")
        plt.axvline(random_boot_vals.median(), color="black", linestyle="--")

        plt.text(
            random_boot_vals.median(), 0, 
            str(random_boot_vals.median().values.round(2)),
            verticalalignment="bottom"
            )

    plt.legend()

def _extend_ds(
    da: xr.DataArray,
    n: int,
    dim: str
    ) -> xr.DataArray:

    extended_da = [da for _ in range(n)]
    return xr.concat(extended_da, dim=dim)
        
    #if dim not in da.dims:
    #    raise ValueError(f"Dimension '{dim}' not found in the data array.")

    ## Find the index of the dimension to repeat along
    #dim_index = list(da.dims).index(dim)

    ## Prepare tiling pattern: replicate n times along the target dimension, 1 along others
    #tile_pattern = [1] * len(da.dims)
    #tile_pattern[dim_index] = n

    ## Tile the data array
    #tiled_data = np.tile(da.values, tuple(tile_pattern))

    #return xr.DataArray(tiled_data, dims=da.dims)

def _xbootstrap(
    x: xr.DataArray,
    y: xr.DataArray,
    func,
    boot_dim: str,
    n_sim: int#,
    #extend_ts: bool = True
    ) -> xr.DataArray:
    
    assert len(x[boot_dim].values) == len(y[boot_dim].values)
    original_len = len(x[boot_dim].values)
    
    #if extend_ts:
    #    x = _extend_ds(x, n_sim, boot_dim)
    #    y = _extend_ds(y, n_sim, boot_dim)
        
    boot_vals = []
    for _ in range(n_sim):
        
        idx = np.random.choice(
            np.arange(len(x[boot_dim].values)), original_len,
            replace=True
            )
        
        boot_vals.append(func(x.isel({boot_dim: idx}), y.isel({boot_dim: idx})))
    
    return xr.concat(boot_vals, dim="n_sim")

    
def _find_crossover_quantile(
    dist1: np.ndarray,
    dist2: np.ndarray,
    quantiles: np.ndarray = np.linspace(0, 1, 100)
    ) -> float:
    
    crossover = []    
    for q in quantiles:
        crossover.append(np.sum(dist1 < q) - np.sum(dist2 < q))
    
    return quantiles[np.argmin(np.abs(crossover))]