# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import torch
torch.manual_seed(123)
import random
random.seed(123)


import torch.nn as nn
import os
import shutil
import itertools


import core.config as config
from chofer_tda_datasets import Reininghaus2014ShrecSynthetic
from chofer_tda_datasets.transforms import Hdf5GroupToDict
from core.utils import *

from torchph.nn.slayer import SLayerExponential, \
SLayerRational, \
LinearRationalStretchedBirthLifeTimeCoordinateTransform, \
prepare_batch, SLayerRationalHat
from sklearn.model_selection import ShuffleSplit
from collections import Counter, defaultdict
from torch.utils.data import DataLoader, SubsetRandomSampler
from collections import OrderedDict
from torch.autograd import Variable
from torch.utils.data import Subset

from sklearn.model_selection import StratifiedShuffleSplit

# %matplotlib notebook

os.environ['CUDA_VISIBLE_DEVICES'] = str(1)

class train_env:
    nu = 0.01
    n_epochs = 300
    lr_initial = 0.5
    momentum = 0.9
    lr_epoch_step = 20
    batch_size = 20
    train_size = 0.9
    minimal_distance_from_diagonal = 0.00001
    

coordinate_transform = LinearRationalStretchedBirthLifeTimeCoordinateTransform(nu=train_env.nu)

def filter_out_diagonal_points(x):
    def fn(dgm):
        i = ((dgm[:, 1] - dgm[:, 0]) > train_env.minimal_distance_from_diagonal).nonzero().squeeze()
        if len(i) == 0:
            return torch.Tensor([])
        else:             
            return dgm[i]
    
    
    return collection_cascade(x, lambda xx: isinstance(xx, torch.Tensor), fn)


dataset = Reininghaus2014ShrecSynthetic(data_root_folder_path=config.paths.data_root_dir)
dataset.data_transforms = [
                           Hdf5GroupToDict(),
                           numpy_to_torch_cascade,   
                           filter_out_diagonal_points,
                           lambda x: collection_cascade(x, 
                                                        lambda x: isinstance(x, torch.Tensor), 
                                                        lambda x: coordinate_transform(x)),
                           ]

dataset.target_transforms = [lambda x: int(x)]


# +
def concat_shrec_sample_target_iter(sample_target_iter):
    x, y = defaultdict(lambda: defaultdict(list)), []
        
    for x_i, y_i in sample_target_iter:
        y.append(y_i)

        for k, v in x_i.items():
            for kk, dgm in v.items():
                x[k][kk].append(dgm)    

    return x, y


class ShrecCollate:   
    def __init__(self, cuda=True):
        self.cuda = cuda
        
    def __call__(self, sample_target_iter):
        x, y = concat_shrec_sample_target_iter(sample_target_iter)

        y = torch.LongTensor(y)    
        
        x = collection_cascade(x,
                           lambda xx: isinstance(xx, list),
                           lambda xx: prepare_batch(xx, 2)) 
        
        if self.cuda:
            # Shifting the necessary parts of the prepared batch to the cuda
            x = collection_cascade(x,
                                   lambda xx: isinstance(xx, tuple),
                                   lambda xx: (xx[0].cuda(), xx[1].cuda(), xx[2], xx[3]))
            y = y.cuda()

        return x, y                       
    
collate_fn = ShrecCollate(cuda=True)


# +
def Slayer(n_elements):
    return SLayerRationalHat(n_elements, radius_init=50, exponent=1)

def LinearCell(n_in, n_out):
    m = nn.Sequential(nn.Linear(n_in, n_out), 
                      nn.BatchNorm1d(n_out), 
                      nn.ReLU(),
                     )
    m.out_features = m[0].out_features
    return m


class ShrecSyntheticModel(nn.Module):
    def __init__(self):
        super().__init__()   
        self.n_elements = 40
        
        self.slayers = ModuleDict()
        for k in (str(i) for i in range(1, 21)):
            self.slayers[k] = ModuleDict()
            for kk in (str(i) for i in range(2)):
                s = Slayer(self.n_elements)
                self.slayers[k][kk] = nn.Sequential(s, nn.BatchNorm1d(self.n_elements))      
            
        cls_in_dim = self.n_elements * 20 * 2
        self.cls = nn.Sequential(
                                LinearCell(cls_in_dim, cls_in_dim),   
                                LinearCell(cls_in_dim, int(cls_in_dim/4)),  
                                nn.Linear(int(cls_in_dim/4), 15))
        
    def forward(self, input):
        x = []
        for k, v in input.items():
            for kk, dgm in v.items():            
                x.append(self.slayers[k][kk](dgm))
        x = torch.cat(x, dim=1)          
        x = self.cls(x)       
                                              
        return x
    
    def center_init(self, sample_target_iter):
        x, _ = concat_shrec_sample_target_iter(sample_target_iter)
        
        x = collection_cascade(x, stop_predicate=lambda e: isinstance(e, list), 
                                  function_to_apply=lambda e: torch.cat(e, dim=0).numpy())        
        
        for k, v in x.items():
            for kk, dgm in v.items():                  
                kmeans = sklearn.cluster.KMeans(n_clusters=self.n_elements, init='k-means++', random_state=123)
                kmeans.fit(dgm)
                centers = kmeans.cluster_centers_
                centers = torch.from_numpy(centers)
                self.slayers[k][kk][0].centers.data = centers
            
            

# +

def experiment():   
    stats_of_runs = []
    
    splitter = StratifiedShuffleSplit(n_splits=10, 
                                      train_size=train_env.train_size, 
                                      test_size=1-train_env.train_size, 
                                      random_state=123)
    train_test_splits = list(splitter.split(X=dataset.targets, y=dataset.targets))
    train_test_splits = [(train_i.tolist(), test_i.tolist()) for train_i, test_i in train_test_splits]
    
    for run_i, (train_i, test_i) in enumerate(train_test_splits):
        print('')
        print('Run', run_i)     

        model = ShrecSyntheticModel()
        model.center_init([dataset[i] for i in train_i])
        model.cuda()

        stats = defaultdict(list)
        stats_of_runs.append(stats)
        
        opt = torch.optim.SGD(model.parameters(), lr=train_env.lr_initial, momentum=train_env.momentum)

        for i_epoch in range(1, train_env.n_epochs+1):      

            model.train()
            
            dl_train = DataLoader(Subset(dataset, train_i),
                                  batch_size=train_env.batch_size, 
                                  collate_fn=collate_fn,
                                  shuffle=True)

            dl_test = DataLoader(Subset(dataset, test_i),
                                 batch_size=train_env.batch_size, 
                                 collate_fn=collate_fn)

            epoch_loss = 0    

            if i_epoch % train_env.lr_epoch_step == 0:
                adapt_lr(opt, lambda lr: lr*0.5)

            for i_batch, (x, y) in enumerate(dl_train, 1):              

                y = torch.autograd.Variable(y)

                def closure():
                    opt.zero_grad()
                    y_hat = model(x)            
                    loss = nn.functional.cross_entropy(y_hat, y)   
                    loss.backward()
                    return loss

                loss = opt.step(closure)

                epoch_loss += float(loss)
                stats['loss_by_batch'].append(float(loss))
                stats['centers'].append(model.slayers['1']['0'][0].centers.data.cpu().numpy())

                print("Epoch {}/{}, Batch {}/{}".format(i_epoch, train_env.n_epochs, i_batch, len(dl_train)), end="       \r")

            stats['train_loss_by_epoch'].append(epoch_loss/len(dl_train))            
                     
            model.eval()    
            true_samples = 0
            seen_samples = 0
            epoch_test_loss = 0
            for i_batch, (x, y) in enumerate(dl_test):

                y_hat = model(x)
                epoch_test_loss += float(nn.functional.cross_entropy(y_hat, torch.autograd.Variable(y.cuda())).data)

                y_hat = y_hat.max(dim=1)[1].data.long()

                true_samples += (y_hat == y).sum().item()
                seen_samples += y.size(0)  

            stats['test_accuracy'].append(true_samples/float(seen_samples))
            stats['test_loss_by_epoch'].append(epoch_test_loss/float(len(dl_test)))
            print('')
            print(true_samples/float(seen_samples))
            
        print('')
        print('acc.', np.mean(stats['test_accuracy'][-10:]))
        return stats_of_runs

stats_of_runs = experiment()
# -

last_10_runs = [np.mean(s['test_accuracy'][-10:]) for s in stats_of_runs]
print(np.mean(last_10_runs))
print(np.std(last_10_runs))
last_10_runs

# +

stats = stats_of_runs[0]
plt.figure()

if 'centers' in stats:
    c_start = stats['centers'][0]
    c_end = stats['centers'][-1]

    plt.plot(c_start[:,0], c_start[:, 1], 'bo', label='center initialization')
    plt.plot(c_end[:,0], c_end[:, 1], 'ro', label='center learned')

    all_centers = numpy.stack(stats['centers'], axis=0)
    for i in range(all_centers.shape[1]):
        points = all_centers[:,i, :]
        plt.plot(points[:, 0], points[:, 1], '-k', alpha=0.25)
        

    plt.legend()
    
plt.figure()
plt.plot(stats['train_loss_by_epoch'], label='train_loss')
plt.plot(stats['test_loss_by_epoch'], label='test_loss')
plt.plot(stats['test_accuracy'], label='test_accuracy')


plt.legend()
plt.show()


StarcoderdataJupyterScriptsDedupFiltered