"""Data generation, preprocessing, and loading."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/source/00_data.ipynb.

# %% auto #0
__all__ = ['AD', 'sbm', 'load_data', 'ConditionalsTransform', 'plot_xf', 'f_a_text', 'plot_xf_D', 'set_plot_xf']

# %% ../nbs/source/00_data.ipynb #df667b48-20b3-4791-9313-1f7adf3bab4e
from fastai.vision.all import *
import andi_datasets
from tqdm.auto import tqdm, trange
from scipy.special import erfcinv

# %% ../nbs/source/00_data.ipynb #c8ea3c80-fe81-4d64-94c9-45a2e8e58fba
from andi_datasets.datasets_theory import datasets_theory
AD = datasets_theory()

# %% ../nbs/source/00_data.ipynb #ed0bf8db-94eb-4eff-86fe-da58581b135c
def sbm(T: int,  # trajectory time steps
        alpha: float,
        sigma: float = 1. )-> np.array:
    '''Creates `T` scaled Brownian motion displacements.'''
    msd = (sigma**2)*np.arange(T+1)**alpha
    dx = np.sqrt(msd[1:]-msd[:-1])
    dx = np.sqrt(2)*dx*erfcinv(2-2*np.random.rand(len(dx)))
    return dx
    # return np.cumsum(dx)-dx[0]  # trajectory centered at origin

# %% ../nbs/source/00_data.ipynb #124738dd-5cd2-44ca-907a-7139201758df
from andi_datasets.utils_trajectories import normalize_fGN

# %% ../nbs/source/00_data.ipynb #b348a472-b55d-460c-87e4-2216320b5551
def load_data(ds_args: dict) -> DataLoaders:
    """Loads a dataset from the given args and returns train and validation `DataLoaders`.
    Trajectories are preprocessed on load and on-the-fly via `ConditionalsTransform`."""

    if ds_args['model']=='fbm':
        ds = AD.create_dataset(T=ds_args["T"], N_models=ds_args["N"],
                               exponents=ds_args["alpha"], models=[2],
                               N_save=ds_args["N_save"],
                               T_save=ds_args["T_save"],
                               save_trajectories=False, load_trajectories=True,
                               path=ds_args["path"])  # trajs, T+2
        ds_labels = ds[:,1][:,None]  # trajs=N*n_alphas,1 # alpha
        n_labels = 1
        ds = ds[:,2:]  # trajs=N*n_alphas, T
        # convert trajectories to displacements:
        ds = np.subtract(ds[:,1:], ds[:,:-1])  # trajs=N*n_alphas, T-1
        # Normalize to the given D
        # take all alphas and assign one D each time
        ds = np.array([normalize_fGN(ds,ds_labels,D,399)
                       for D in ds_args['D']
                      ]).reshape(-1,ds_args["T"]-1)
        ds_labels =np.concatenate((np.tile(ds_labels, (len(ds_args['D']),1)),
                                   np.array(ds_args['D'])[:,None].repeat(len(ds_labels),axis=0)
                                  ),axis=-1)
    elif ds_args['model']=='sbm':
        assert ds_args["N_save"]>=ds_args["N"]
        ds_labels = np.empty((ds_args["N"]*len(ds_args["alpha"])*len(ds_args["D"]),2))
        ds        = np.empty((ds_args["N"]*len(ds_args["alpha"])*len(ds_args["D"]),ds_args["T"]-1))
        n = 0
        f = np.load(ds_args["path"]+'sbm.npz')
        for i,a in enumerate(ds_args["alpha"]):
            for j,D in enumerate(ds_args["D"]):
                k = f'{a:.3g}'+f',{D:.3g}'
                ds_labels[n:n+ds_args["N"]] = f[k][:ds_args["N"],:2]
                ds[n:n+ds_args["N"]] = f[k][:ds_args["N"],2:2+ds_args["T"]-1]
                n =  n+ds_args["N"]
    # trajs=N*n_alphas, 2
    n_labels=2
    # ensure data type
    ds = ds.astype(np.float32);  ds_labels = ds_labels.astype(np.float32)
    ds = np.concatenate((ds_labels,ds),axis=-1)
    tfm = ConditionalsTransform(n_labels=n_labels)

    n_inp = 1
    splits_idx = RandomSplitter(valid_pct=ds_args['valid_pct'], seed=ds_args['seed'])(np.arange(len(ds)))
    tfl = TfmdLists(ds, tfm, splits=splits_idx,)  # iter
    return tfl.dataloaders(ds_args['bs'], n_inp=n_inp, num_workers=2, pin_memory=True, drop_last=True)

class ConditionalsTransform(Transform):
    """Extracts the positional information from the input array,
    inserts a channel dimension for convolutions, and discards the labels."""
    def __init__(self, n_labels=1): store_attr()
    def encodes(self, andi_dataset_item):
        ds,ds_labels = andi_dataset_item[self.n_labels:], andi_dataset_item[:self.n_labels]
        # for convolutions insert a channel dimension.
        ds_ = np.expand_dims(ds, axis=-1)
        channel_dim = 0
        ds = np.expand_dims(ds, axis=channel_dim)
        return (ds, ds_)

# %% ../nbs/source/00_data.ipynb #20ddfd63-51d2-467d-8acc-f84f046651a5
def plot_xf(x, f, ds_args:dict, intersect_idx,
            x_label:str='', y_label:str='', title:str='',
            x_scale:str='linear',y_scale:str='linear',):
    """Plot a set of input arrays `x` transformed by `f`
    and colored by a parameter combination."""
    for i,a in enumerate(ds_args["alpha"]):
        for j,D in enumerate(ds_args["D"]):
            plt.plot(f(x[intersect_idx[i,j]]), label=f'{a:.3g}, {D:.3g}');
    plt.xlabel(x_label); plt.ylabel(y_label, rotation=0);
    plt.xscale(x_scale); plt.yscale(y_scale); plt.title(title);

def f_a_text(a,y):
    """Format an anomalous exponent `a` value as a text annotation in (0.95,y) of axes coordinates"""
    plt.text(0.95,y[-1],f'{a:.3g}',
             transform=matplotlib.transforms.blended_transform_factory(plt.gca().transAxes, plt.gca().transData));

def plot_xf_D(x,f, ds_args:dict, intersect_idx,
              f_a=None,f_D=None,f_aD=None,each_D=True,each_g=True, **kwargs):
    """Plots `f(x)` using the provided indices and auxiliary functions"""
    kwargs.update(dict(
        legend_title= r'$\alpha, D$',
        ncol=len(ds_args["alpha"])*len(ds_args["D"])//8 +1,
    ))
    for j,D in enumerate(ds_args["D"]):
        if each_D: plt.figure()
        for i,a in enumerate(ds_args["alpha"]):
            p = plt.plot(f(x[intersect_idx[i,j]]), label=f'{a:.2g}, {D:.2g}');
            if f_a: f_a(a,p[0]._y)
            if f_aD: f_aD(a,D)
        if f_D: f_D(D)
        if each_D:
            if 'save' in kwargs: kwargs['save']=kwargs['save'][:kwargs['save'].rfind('_D')]+f'_D{D:.3g}'.replace('.','')
            set_plot_xf(**kwargs)
    if not each_D: set_plot_xf(**kwargs)

def set_plot_xf(x_label:str='', y_label:str='', title:str='',
                x_scale:str='linear', y_scale:str='linear',
                legend_title:str='', save:bool=False,
                ncol:int=1, show_legend:bool=True):
    """Configure axis labels and formatting for trajectory plots."""
    plt.xlabel(x_label);plt.ylabel(y_label, rotation=0, ha='right');
    plt.xscale(x_scale);plt.yscale(y_scale);
    if ncol<4 and  show_legend==True: plt.legend(title=legend_title, bbox_to_anchor=(1,1), ncol=ncol);
    plt.title(title);plt.grid();
    if save:
        plt.savefig(save +'.png', bbox_inches='tight', pad_inches=0.1);
        print('Saved:' + save +'.png')
        plt.close();
