#| default_exp aa_mixer_toy

#| export
import math
import os
import matplotlib.pyplot as plt
from PIL import Image
import random
import string
import numpy as np

import torch
from torch import optim, nn, Tensor
from torch.nn import functional as F
from torch.utils import data as torchdata

from einops import rearrange
from tqdm.auto import tqdm
import wandb
from aeiou.core import get_device

train_dl = torchdata.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_dl = torchdata.DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
val_iter = iter(val_dl)

#| export
def get_demo_batch(val_iter, 
                   demo_line = True, 
                   debug=False):
    demo_batch = next(val_iter).clone()
    absmax = (torch.abs(demo_batch)).max().item()

    if demo_line:  # make demo batch lie along a line in order to showcase nonlinearity
        demo_batch[:,1] = demo_batch[:,0]  # line y = x
        demo_batch[::2,0] *= -1            # even-numbered points go to y = -x
        demo_batch[::4,0] = 0.01   # data along y axis
        demo_batch[1::4,1] = 0.01  # data along x axis
        demo_batch[::5] = absmax*(2*torch.rand((demo_batch[::5]).shape)-1)   # make some points random
        demo_batch[1::7] = absmax*(2*torch.rand((demo_batch[1::7]).shape)-1)  # and more points random
    return demo_batch

#| export 
class EmbedBlock(nn.Module):
    def __init__(self, in_dims:int, out_dims:int, act=nn.GELU(), resid=True, use_bn=False, requires_grad=True, **kwargs) -> None:
        "generic little block for embedding stuff.  note residual-or-not doesn't seem to make a huge difference for a-a"
        super().__init__()
        self.in_dims, self.out_dims, self.act, self.resid = in_dims, out_dims, act, resid
        self.lin = nn.Linear(in_dims, out_dims, **kwargs)
        self.bn = nn.BatchNorm1d(out_dims) if use_bn else None # even though points in 2d, only one non-batch dim in data

        if requires_grad == False:
            self.lin.weight.requires_grad = False
            self.lin.bias.requires_grad = False

    def forward(self, xin: Tensor) -> Tensor:
        x = self.lin(xin)
        if self.act is not None: x = self.act(x)
        if self.bn is not None: x = self.bn(x)   # re. "BN before or after Activation? cf. https://github.com/ducha-aiki/caffenet-benchmark/blob/master/batchnorm.md"
        return xin + x if (self.resid and self.in_dims==self.out_dims) else x 

# test that
emb_test = EmbedBlock(2, 2, use_bn=True, requires_grad=False)
#res_test = ResBlock(2, 2)

#| export
def friendly_tanh(x, a=3, b=0.1, s=0.9):
    "like a tanh but not as harsh"
    return (torch.tanh(a*x) + b*x)*s
    #return torch.tanh(a*x) + b*torch.sgn(x)*torch.sqrt(torch.abs(x))
    
def compressor(x, thresh=0.7, ratio=10):
    "another nonlinearity: kind of like an audio compressor"
    return torch.where(torch.abs(x) < thresh, x, torch.sgn(x)*(1/ratio*(torch.abs(x)-thresh)+thresh))

#| export    
class TwistAndScrunch(nn.Module):
    """for now just something random with a bit of nonlinearity
    """
    def __init__(self,
        in_dims=2,        # number of input dimensions to use
        emb_dims=2,       # number of output dimensions to use
        hidden_dims=32,   # number of hidden dimensions
        scrunch_fac=1.7, twist_fac=0.2, debug=False, act=nn.GELU(),
        ): 
        super().__init__()
        self.sf, self.debug, = scrunch_fac, debug, 
        self.pi = torch.acos(torch.zeros(1)).item() * 2
        
        self.decoder = nn.Sequential(  
            EmbedBlock(emb_dims, hidden_dims, act=act),
            EmbedBlock(hidden_dims, hidden_dims, act=act),            
            nn.Linear(hidden_dims, in_dims) )        
        
        # note we'll never use this encoder. Just included this to keep other codes from crashing. ;-) 
        self.encoder = nn.Sequential(   # note that this is not going to be used.
            EmbedBlock(emb_dims, hidden_dims, act=act),
            EmbedBlock(hidden_dims, hidden_dims, act=act),            
            nn.Linear(hidden_dims, in_dims)
        ) 
            
    def encode(self, v):  # some nonlinear function i made up
        vp = v.clone()    
        vp[:,0], vp[:,1] = vp[:,0] - 0.5*friendly_tanh(v[:,1],a=3), vp[:,1] + 0.7*friendly_tanh(v[:,0],a=1.2)
        return friendly_tanh(self.sf*vp,a=2,b=0.1)
    
    def decode(self, emb, exact=False):
        return self.decoder(emb)
        
    def forward(self, x, debug=False, exact=False):
        debug_save, self.debug = self.debug, debug  # only turn on debug for this one call
        emb = self.encode(x)
        outs =  self.decode(emb, exact=exact), emb
        self.debug = debug_save
        return outs

torch.manual_seed(seed)
#given_model = SimpleAutoEncoder(in_dims=in_dims, emb_dims=emb_dims).to(device)
given_model = TwistAndScrunch()

given_model = given_model.to(device)

#| export
def viz_given_batch(
    val_batch,            # should be from val_batch
    given_model,         # the main autoencoder model we are splicing into
    aa_model=None,       # our audio algebra model
    debug=False, 
    device=get_device(),
    return_img = True,
    show_recon=False,
    ):
    markers=['o', '+', 'x', '*', '.', 'X', "v", "^", "<", ">", "8", "s", "p", "P", "h", "H", "D", "d","|","_"]
    colors=['blue', 'magenta','orange', 'green', 'red', 'cyan', 'black']
    markersize=2.5
    
    val_batch = val_batch.to(device)
    
    with torch.no_grad():
        emb_batch = given_model.encode(val_batch)
        recon_batch = given_model.decode(emb_batch)
        
    fig, ax = plt.subplots(1,2, figsize=(11,5))
    ax[0].set_title(f"input space.  $x \in X$")
    ax[1].set_title(f"given encoder space. $y \in Y$")
    
    if debug: 
        print("absmax of val_batch = ",torch.abs(val_batch).max().item())
        print("absmax of emb_batch = ",torch.abs(emb_batch).max().item())
    
    nm = len(markers)
    for mi in range(nm):
    #for i in range(val_batch.shape[0]):  # TODO: looping over every points is very slow, but apparently needed for arrows
        for bi,b in enumerate([val_batch, emb_batch, recon_batch]):
            x, y = b[mi::nm,0].cpu().numpy(), b[mi::nm,1].cpu().numpy()
            if (bi<2) or show_recon: ax[bi%2].plot(x, y, marker=markers[mi], markersize=markersize, color=colors[bi], linestyle='None', label='' if mi > 0 else ['input', 'embedding', 'recon'][bi])
    ax[0].legend(loc='upper left')
    val_plot_fname = f'{SCRATCH_DIR}/val_plot{RUN_SUFFIX}{random.randint(1,5)}.png'
    plt.savefig(val_plot_fname, bbox_inches='tight', dpi=300)
    plt.close()
    im = Image.open(val_plot_fname)
    return im

#| export
class AudioAlgebra(nn.Module):
    """
    Main AudioAlgebra model. Contrast to aa-mixer code, keep this one simple & move mixing stuff outside
    """
    def __init__(self, dims=2, hidden_dims=64, act=nn.GELU(), use_bn=False, resid=True, block=EmbedBlock):
        super().__init__()
        self.encoder = nn.Sequential(  
            block( dims,        hidden_dims, act=act,  use_bn=use_bn, resid=resid),
            block( hidden_dims, hidden_dims, act=act,  use_bn=use_bn, resid=resid),
            block( hidden_dims, hidden_dims, act=act,  use_bn=use_bn, resid=resid),
            block( hidden_dims, dims,        act=None, use_bn=use_bn, resid=resid),
        )
        self.decoder = nn.Sequential(  # same as encoder, in fact. 
            block( dims,        hidden_dims, act=act,  use_bn=use_bn, resid=resid),
            block( hidden_dims, hidden_dims, act=act,  use_bn=use_bn, resid=resid),
            block( hidden_dims, hidden_dims, act=act,  use_bn=use_bn, resid=resid),   
            block( hidden_dims, dims,        act=None, use_bn=use_bn, resid=resid),
        )
            
    def encode(self,x):
        return x + self.encoder(x) 

    def decode(self,x):
        return x + self.decoder(x)

    def forward(self, 
        x   # the enbedding vector from the given encoder
        ):
        xprime = self.encode(x)
        xprimeprime = self.decode(xprime)  # train system to invert itself (and hope it doesn't all collapse to nothing!)
        return xprime, xprimeprime  # encoder output,  decoder output

#| export 
def get_stems_faders(batch, #  "1 stem" (or batch thereof) already drawn fron the dataloader (val or train)
                     dl_iter,  # pre-made the iterator for the/a dataloader
                     dl,       # the dataloader itself, for restarting
                     maxstems=2,  # how many total stems will be used, i.e. draw maxstems-1 new stems from dl_iter
                     unity_gain=False,  # this will force all faders to be +/-1 instead of random numers
                     debug=False):
    "grab some more inputs and multiplies and some gain values to go with them"
    nstems = random.randint(2, maxstems)
    if debug: print("maxstems, nstems =",maxstems, nstems)
    device=batch.device
    faders = torch.sgn(2*torch.rand(nstems)-1)  # random +/- 1's
    if not unity_gain:
        faders += 0.5*torch.tanh(2*(2*torch.rand(nstems)-1))  # gain is now between 0.5 and 1.5
    stems = [batch]                  # note that stems is a list
    for i in range(nstems-1):        # in addtion to the batch of stem passed in, grab some more
        try: 
            next_stem = next(dl_iter).to(device)    # this is just another batch of input data
        except StopIteration:
            dl_iter = iter(dl)       # time to restart. hoping this propagates out as a pointer
            next_stem = next(dl_iter).to(device)
        if debug: print("  next_stem.shape = ",next_stem.shape)
        stems.append(next_stem)
    return stems, faders.to(device), dl_iter  # also return the iterator

def do_mixing(stems, faders, given_model, aa_model, device, debug=False, **kwargs):
    """
    here we actually encode inputs.
    0's denote values in the given model space, non-0's denode those in our aa_model
    e.g., "y" denotes an embedding from the frozen encoder, "z" denotes re-mapped embeddings

    """
    zs, ys, zsum, ysum, yrecon_sum, fadedstems = [], [], None, None, None, []
    mix = torch.zeros_like(stems[0]).to(device)
    #if debug: print("do_mixing: stems, faders =",stems, faders)
    for s, f in zip(stems, faders):   # iterate through list of stems, encode a bunch of stems at different fader settings
        fadedstem = (s * f).to(device)                 # audio stem adjusted by gain fader f
        with torch.no_grad():
            y = given_model.encode(fadedstem)  # encode the stem
        z, y_recon = aa_model(y)             # <-- this is the main work of the model
        zsum = z if zsum is None else zsum + z # <---- compute the sum of all the z's so far. we'll end up using this in our (metric) loss as "pred"

        mix += fadedstem                 # make full mix in input space
        with torch.no_grad():
            ymix = given_model.encode(mix)  # encode the mix in the given model
        zmix, ymix_recon = aa_model(ymix)   #  <----- map that according to our learned re-embedding. this will be the "target" in the metric loss

        # for diagnostics:
        ysum = y if ysum is None else ysum + y   # = sum of embeddings in original model space; we don't really care about ysum except for diagnostics
        yrecon_sum = y_recon if yrecon_sum is None else yrecon_sum + y_recon   # = sum of embeddings in original model space; we don't really care about ysum except for diagnostics

        zs.append(z)              # save a list of individual z's
        ys.append(y)            # save a list of individual y's
        fadedstems.append(fadedstem) # safe a list of each thing that went into the mix
        
    archive = {'zs':zs, 'mix':mix, 'znegsum':None, 'ys': ys, 'ysum':ysum, 'ymix':ymix, 'ymix_recon':ymix_recon, 'yrecon_sum':yrecon_sum, 'fadedstems':fadedstems} # more info for diagnostics

    return zsum, zmix, archive  # we will try to get these two to be close to each other via loss. archive is for diagnostics

aa_use_bn = False  # batch norm? 
aa_use_resid = True # use residual connections? (doesn't make much difference tbh)
hidden_dims = 64   # number of hidden dimensions in aa model. usually was 64

torch.manual_seed(seed) # chose this value because it shows of nice nonlinearity
aa_model = AudioAlgebra(dims=emb_dims, hidden_dims=hidden_dims, use_bn=aa_use_bn, resid=aa_use_resid).to(device) 

#| export
def multimarkerplot(ax, data, color='blue', markersize=2.7, label=''):
    "little thingy for using lots of symbols for a plot"
    markers = ['o', '+', 'x', '*', '.', 'X', "v", "^", "<", ">", "8", "s", "p", "P", "h", "H", "D", "d","|","_"]
    nm = len(markers)
    for mi in range(nm):
        ax.plot(data[mi::nm,0], data[mi::nm,1], marker=markers[mi], markersize=markersize, color=color, linestyle='None', label=label if mi==0 else '') 

def viz_aa_demo(
    zsum, zmix,
    archive,
    aa_model,
    debug=False,
    ):
    "for visualizing what aa model is doing"
    colors=['blue', 'red', 'green', 'magenta', 'orange', 'cyan', 'black',]

    col_labels = ["inputs", "given representations", "a-a embeddings"]
    nrows, ncols = 2, len(col_labels)
    figsize = (4*ncols, 4*nrows)
    fig, ax = plt.subplots(nrows,ncols, figsize=figsize)
    
    # top row: stems*faders,  ys,  zs:  each of these is a list
    keys, varnames = ['fadedstems', 'ys', 'zs'], ['x', 'y', 'z']    
    for coli, key in enumerate(keys):
        ax[0,coli].set_title(f"{col_labels[coli]} ${varnames[coli]}$")
        datalist = [x.detach().cpu().numpy() for x in archive[key]]
        for i, stem in enumerate(datalist):
            label = f"{varnames[coli]}{i+1}"
            if coli > 0:
                label += f" := f(x{i+1})" if coli==1 else f" := h(y{i+1})"
            multimarkerplot( ax[0,coli], stem, color=colors[i], label=label)
        ax[0,coli].legend(fancybox=True, framealpha=0.6, loc='upper left', prop={'size': 9})

    # bottom row: mix
    archive["zmix"] = zmix  # lil prep
    keys = ['mix', 'ymix', 'zmix']
    annotations = [' := x1 + x2', r' := $f\,$(mix)', r' := $h$(ymix)']
    for coli, key in enumerate(keys):
        ax[1,coli].set_title(f"{col_labels[coli]}")
        data = archive[key].detach().cpu().numpy()
        multimarkerplot( ax[1,coli], data, label=key+annotations[coli], color=colors[2])
        if key=='ymix': # add another plot
            multimarkerplot( ax[1,coli], archive["ymix_recon"].detach().cpu().numpy(), label=r'ymix_recon := $h^{-1}$(zmix)', color=colors[4])
        elif key=='zmix': 
            multimarkerplot( ax[1,coli], zsum.detach().cpu().numpy(), label='zsum := z1 + z2', color=colors[3])
        ax[1,coli].legend(fancybox=True, framealpha=0.6, loc='upper left', prop={'size': 9})
        
    if debug: print("saving to file")
    filename = f'{SCRATCH_DIR}/aa_val_plot{RUN_SUFFIX}.png'
    plt.savefig(filename, bbox_inches='tight', dpi=225)
    plt.close()
    if debug: print("loading file")
    im = Image.open(filename)
    return im

mseloss = nn.MSELoss()

def rel_loss(y_pred: torch.Tensor, y: torch.Tensor, eps=1e-3) -> float:
    "relative error loss   --- note we're never going to actually use this. it was just part of development"
    e = torch.abs(y.view_as(y_pred) - y_pred) / ( torch.abs(y.view_as(y_pred)) + eps ) 
    return torch.median(e)

#| export
def vicreg_var_loss(z, gamma=1, eps=1e-4):
    std_z = torch.sqrt(z.var(dim=0) + eps)
    return torch.mean(F.relu(gamma - std_z))   # the relu gets us the max(0, ...)

#| export
def off_diagonal(x):
    n, m = x.shape
    assert n == m
    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()

def vicreg_cov_loss(z):
    "the regularization term that is the sum of the off-diagaonal terms of the covariance matrix"
    cov_z = torch.cov(z.T)   
    return off_diagonal(cov_z).pow_(2).sum().div(z.shape[1])

#| export
model_path = 'aa_model.pth'

def load_aa_checkpoint(model_path='aa_model.pth', emb_dims=2, device=get_device()):
    aa_model = None
    try:
        aa_model = torch.load(model_path, map_location=torch.device('cpu')).to(device)
    except Exception as e:
        print(f"Error trying to load presaved models:\n{e}\n\nStarting from scratch.")
        return None, None
    print("Successfully loaded pretrained models")
    return aa_model


def save_aa_checkpoint(aa_model, model_path='aa_model.pth', suffix=''):
    "save a model checkpoint -- full version. note we save to scratch so as not to overwrite existing file"
    torch.save(aa_model, SCRATCH_DIR+'/'+model_path.replace(".pth", f"{suffix}.pth"))        

#| export
train_new_aa = True  # set to false to disable this training


def sayno(b, suffix=''): # just a little util for naming runs
    return f'no{suffix}' if (b is None) or (b==False) else ''


def train_aa_vicreg(debug = False): 
    "packaging as def to re-use later.  this will use global variables, not sorry"

    max_epochs = 40
    lossinfo_every, viz_demo_every =   20, 1000   # in units of steps
    checkpoint_every = 10000
    max_lr= 0.002
    
    # optimizer and learning rate scheduler
    opt = optim.Adam([*aa_model.parameters()], lr=5e-4) 
    total_steps = len(train_dataset)//batch_size * max_epochs
    print("total_steps =",total_steps)  # for when I'm checking wandb
    scheduler = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=max_lr, total_steps=total_steps)

    # rename our wandb run with some significant info
    wandb.init(project='aa-toy-vicreg')
    run_name = wandb.run.name   
    wandb.run.name = f"h{hidden_dims}_{'bn_' if aa_use_bn else ''}bs{batch_size}_lr{max_lr}{RUN_SUFFIX}"
    print("New run name =",wandb.run.name)
    wandb.run.save()    
    
    # training loop
    train_iter = iter(train_dl) # this is only for use with get_stems_faders
    epoch, step = 0, 0
    torch.manual_seed(seed) # for reproducibility
    while (epoch < max_epochs) or (max_epochs < 0):  # training loop
        with tqdm(train_dl, unit="batch") as tepoch:
            for batch in tepoch:   # train
                opt.zero_grad()
                log_dict = {}
                batch = (1.2*(2*torch.rand(batch.shape, device=device)-1)) # overwrite batch with unending random data!
                #batch = batch.to(device)

                stems, faders, train_iter = get_stems_faders(batch, train_iter, train_dl, debug=debug)

                # vicreg: 1. invariance
                zsum, zmix, archive = do_mixing(stems, faders, given_model, aa_model, device, debug=debug)
                mix_loss = mseloss(zsum, zmix)  

                var_loss = (vicreg_var_loss(zsum) + vicreg_var_loss(zmix))/2    # vicreg: 2. variance
                cov_loss = (vicreg_cov_loss(zsum) + vicreg_cov_loss(zmix))/2    # vicreg: 3. covariance


                # reconstruction loss: inversion of aa map h^{-1}(z): z -> y,  i.e. train the aa decoder
                y = given_model.encode(batch)
                z, yrecon = aa_model.forward(y)       # (re)compute ys for one input batch (not stems&faders)
                aa_recon_loss = mseloss(y, yrecon)     
                aa_recon_loss = aa_recon_loss + mseloss(archive['ymix'], archive['ymix_recon'])  # also recon of the  mix ecoding
                #aa_recon_loss  = aa_recon_loss +  mseloss(archive['ysum'], archive['yrecon_sum']) # Never use this:  ysum shouldn't matter / is poorly defined
           
                loss = mix_loss + var_loss + cov_loss + aa_recon_loss     # --- full loss function
                
                log_dict['train_loss'] = loss.detach()                    # --- this is the full loss 
                log_dict['mix_loss'] = mix_loss.detach() 
                log_dict['aa_recon_loss'] = aa_recon_loss.detach()
                log_dict['var_loss'] = var_loss.detach() 
                log_dict['cov_loss'] = cov_loss.detach() 
                log_dict['learning_rate'] = opt.param_groups[0]['lr']
                log_dict['epoch'] = epoch

                if step % lossinfo_every == 0: 
                    tepoch.set_description(f"Epoch {epoch+1}/{max_epochs}")
                    tepoch.set_postfix(loss=loss.item())         #  TODO: use EMA for loss display? 

                loss.backward()
                opt.step()  

                with torch.no_grad():
                    # run on demo data (same each time)
                    demo_zsum, demo_zmix, demo_archive = do_mixing(demo_stems, demo_faders, given_model, aa_model, device, debug=debug)
                    demo_mix_loss = mseloss(demo_zsum, demo_zmix)
                    log_dict['demo_mix_loss'] = demo_mix_loss.detach()
                    
                    if step % viz_demo_every == 0:   # but only draw the viz every now and then. 
                        try:     # to avoid "OSError: image file is truncated"
                            im = viz_aa_demo(demo_zsum, demo_zmix, demo_archive, aa_model)
                            log_dict["aa_mapping"] = wandb.Image(im)
                        except:
                            pass

                if step % checkpoint_every == 0:
                    save_aa_checkpoint(aa_model, suffix=RUN_SUFFIX+f"_{step}")

                wandb.log(log_dict)

                scheduler.step()   
                step += 1

        epoch += 1
    #----- training loop finished
    
    save_aa_checkpoint(aa_model, suffix=RUN_SUFFIX+f"_{step}")
    

#| export
def inp_to_aa(x):
    "z = h(f(x))"
    with torch.no_grad():
        y = given_model.encode(x.to(device))
        return aa_model.encode(y),  y

def aa_to_inp(z):
    "x = h^-1 ( f^-1 ( z ) )"
    with torch.no_grad():
        y = aa_model.decode(z.to(device))
        return given_model.decode(y), y
    
    
def plot_kmwq_demo(q, forward=True, labels = ['man', 'king','woman','queen']):  
    # forward = inputs to z's, backaward = z's to inputs
    " little routine to plot our math ops"
    if forward:
        xs = q
        zs, ys = inp_to_aa(q)
        maybe = zs[1] - zs[0] + zs[2]
    else:
        zs = q
        xs, ys = aa_to_inp(q)
        maybe = xs[1] - xs[0] + xs[2]
        
    [xs, zs, ys, maybe] = [ q.cpu().numpy() for q in [xs, zs, ys, maybe]]

    fig, ax = plt.subplots(1,3,figsize=(12,4))
    
    titles, qs  = ['input','given emb','aa emb'], [xs, ys, zs]
    for i, q in enumerate(qs):
        ax[i].set_title(titles[i])
        ax[i].plot(q[:,0], q[:,1], marker='o') 
   
    print("orange dot shows the 'guess'")
    ax[forward*2].plot(maybe[0], maybe[1], marker="o") 
    
    for i,lab in enumerate(labels): # show the individual data points and label them
        ax[0].text(xs[i,0],xs[i,1], lab)
        ax[2].text(zs[i,0],zs[i,1], lab)

    points = np.vstack((xs,zs))
    valmin, valmax = 1.1*points.min(axis=0), 1.1*points.max(axis=0)
    for i in range(3):
        ax[i].set_xlim(-2,2)
        ax[i].set_ylim(-2,2)

    plt.plot()