# ---
# 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 matplotlib
matplotlib.use('Agg')

# +
from tqdm import tqdm
tqdm.monitor_interval = 0

from fastai.imports import *
from fastai.transforms import *
from fastai.conv_learner import *
from fastai.model import *
from fastai.dataset import *
from fastai.sgdr import *
from fastai.plots import *
import json
import pandas as pd
from sklearn.metrics import *
# -

# %matplotlib inline
# %reload_ext autoreload
# %autoreload 2

from mobile_net import *

from cifar10 import get_data

PATH = Path("/scratch/arka/miniconda3/external/fastai/courses/dl2/data/cifar10/")
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

sz = 32
bs = 128

data = get_data(sz, bs, PATH)


class SiameseDS(BaseDataset):
    def __init__(self, ds, wgt_pos_neg=0.5):
        self.ds = ds
        super().__init__(ds.transform)
        self.label_set = set(data.trn_ds.y)
        self.l2i = {l: np.where(l == ds.y)[0] 
                    for l in self.label_set}
        self.wgt = T(wgt_pos_neg)
    
    def __len__(self):
        return self.get_n()
        
    def get_n(self):
        return self.ds.n
    
    def get_c(self):
        return 2
    
    def get_sz(self):
        return self.ds.sz
    
    def get_x(self, i):
        return self.ds.get_x(i)
    
    def get_y(self, i):
        return self.ds.get_y(i)
    
    def __getitem__(self, idx):
        targ = np.random.randint(0, 2)
        x1, y1 = self.get1item(idx)
        if targ == 1:
            new_idx = np.random.choice(self.l2i[y1])
            x2, y2 = self.get1item(new_idx)
            assert y1 == y2
        else:
            new_c = np.random.choice(sorted(list(self.label_set - {y1})))
            new_idx = np.random.choice(self.l2i[new_c])
            x2, y2 = self.get1item(new_idx)
            assert y1 != y2
        return [x1, x2, y1, y2, targ, targ]


class SiameseData(ImageData):
    @classmethod
    def from_image_classifier_data(cls, data, bs=64, num_workers=4):
        trn_ds = SiameseDS(data.trn_dl.dataset)
        val_ds = SiameseDS(data.val_dl.dataset)
        fix_ds = SiameseDS(data.fix_dl.dataset)
        aug_ds = SiameseDS(data.aug_dl.dataset)
        res = [trn_ds, val_ds, fix_ds, aug_ds]
        if data.test_dl is not None:
            test_ds = SiameseDS(data.test_dl.dataset)
            test_aug = SiameseDS(data.test_aug_dl.dataset)
            res += [test_ds, test_aug]
        else:
            res += [None, None]
        return cls(data.path, res, bs, num_workers, classes=['pos', 'neg'])


# +
class SiameseModel(nn.Module):
    def __init__(self, mdl):
        super().__init__()
        self.mdl = mdl
        self.pdist = nn.PairwiseDistance()
        self.ffc = nn.Linear(1, 2)
    
#     def emb_out(self, inp):
#         out = F.relu(self.mdl.bn1(self.mdl.conv1(inp)))
#         out = F.adaptive_avg_pool2d(self.mdl.lyrs(out), 1)
#         out = out.view(out.size(0), -1)
#         return out
    
    def forward(self, inp0, inp1, y0, y1, targ):
        o0 = self.mdl.emb_out(inp0)
        o1 = self.mdl.emb_out(inp1)
#         import pdb; pdb.set_trace()
        out = self.pdist(o0, o1)
        out = F.log_softmax(self.ffc(out),dim=-1)
#         import pdb; pdb.set_trace()
        return out


# -

class ContrastiveLoss(nn.Module):
    def __init__(self, margin):
        super().__init__()
        self.margin = margin
    def forward(self, outp, targ, size_average=True):
        tmp1 = targ.float() * outp**2 
        tmp2 = (1 - targ.float())
        tmp3 = F.relu(self.margin - outp)
        tmp4 = tmp2 * tmp3**2
        res = (tmp1 + tmp4) / 2
        return res.mean() if size_average else res.sum()


def acc_siamese(preds, targs):
    s = np.zeros(targs.shape)
    p = np.where(preds > 0.5)[0]
    s[p] = 1
    return (s==targs.cpu().numpy()).mean()


md_mbl = mblnetv1(depthwise_block, 
              inc_list=[64, 64, 128, 256], 
              inc_scale = 1, 
              num_blocks_list=[2, 2, 2], 
              stride_list=[1, 2, 2], 
              num_classes=10)

learn1 = ConvLearner.from_model_data(md_mbl, data)
learn1.load('mobilenetv1_3')

learn1.fit(0, 1)

# md_mbl = EmbeddingNet()
# md_mbl_trained = learn1.model
sia_mdl = SiameseModel(md_mbl)
sia_data = SiameseData.from_image_classifier_data(data)
learn = ConvLearner.from_model_data(sia_mdl, sia_data)

# +
# learn.crit = ContrastiveLoss(margin=1)
# learn.predict()
# learn.fit(1e-2, 3, cycle_len=1, metrics=[acc_siamese])
# -

learn.fit(1e-1, 3, cycle_len=1)

learn.fit(1e-1, 2, cycle_len=10)

learn.fit(1e-1, 2, cycle_len=10)

learn.fit(1e-1, 2, cycle_len=10)

learn.fit(1e-1, 2, cycle_len=20)

learn.save('mbnet1_siam1')

learn.load('mbnet1_siam1')

learn.fit(1e-2, 2, cycle_len=1, best_save_name='mbnet_siam2')

learn.load('mbnet_siam2')
learn.unfreeze()

learn.fit(1e-1, 2, cycle_len=1, use_clr=(20, 5), best_save_name='mbnet_siam3')

StarcoderdataJupyterScriptsDedupFiltered