from fastai2.basics import *
from fastai2.callback.all import *
from fastai2.vision.all import *
from torch import torch

from fa_convnav.navigator import *
from pandas import DataFrame

pets = DataBlock(blocks=(ImageBlock, CategoryBlock), 
                 get_items=get_image_files, 
                 splitter=RandomSplitter(),
                 get_y=RegexLabeller(pat = r'/([^/]+)_\d+.jpg$'),
                 item_tfms=Resize(460),
                 batch_tfms=[*aug_transforms(size=224, max_rotate=30, min_scale=0.75), Normalize.from_stats(*imagenet_stats)])

dls = pets.dataloaders(untar_data(URLs.PETS)/"images",  bs=128)

%%capture
def run_tests(cn_test, i):
  test_df = cn_test._cndf

  test_eq(type(cn_test._cndf), res[0]) 
  test_eq(len(cn_test._cndf), res[1])             # rows
  test_eq(len(cn_test._cndf.columns), res[2])     # columns   

  test_df['lyr_obj'] = None
  test_eq(len(cndf_search(test_df, 12)), res[3])
  test_eq(len(cndf_search(test_df, ['0.6.1.conv2', '0.0.6', '0.0.6', '0.0.6', '0.0.4', '0.6'][i])), res[4])
  test_eq(len(cndf_search(test_df, ['0.6', '0.0.6', '0.0.6', '0.0.6', '0.0.4.2', '0.6'][i], exact=False)), res[5])
  test_eq(len(cndf_search(test_df, [{'Module_name': '0.6', 'Layer_description':'Conv2d'}, \
                                    {'Module_name': '1.0', 'Container_child':'AdaptiveConcatPool2d'}, \
                                    {'Module_name': '1.0', 'Container_child':'AdaptiveConcatPool2d'}, \
                                    {'Module_name': '0.0.6', 'Layer_description':'Conv2d'}, \
                                    {'Module_name': '0.0.4.2', 'Layer_description':'Conv2d'}, \
                                    {'Module_name': '0.6', 'Layer_description':'Conv2d'}, \
                                    ][i], exact=True)), res[6])
  test_eq(len(cndf_search(test_df, ['0.6', '0.5'], exact=False)), res[7])
  test_eq(cndf_search(test_df, ('0.6', '0.5'), exact=False), res[8])

  cn_test.view()
  cn_test.head
  cn_test.body
  cn_test.divs 
  test_eq(len(cn_test.linear_layers), res[9])                                  
  test_eq(len(cn_test.dim_transitions), res[10])
  test_eq(len(cn_test.find_block('0.4.1')), res[11])
  test_eq(len(cn_test.find_block('0.4.1', layers=False)), res[12])
  test_eq(len(cn_test.find_conv('first', 5)), res[13]) #revise to 5 ater importing chnages from core.ipynb+ below 3 -> 5 
  test_eq(len(cn_test.children), res[14])
  test_eq(len(cn_test.blocks), res[15])
  test_eq(len(cn_test.spread('conv', 8)), res[16])
  del(cn_test)

models_to_test = [
                  ('resnet18',      resnet18,       [DataFrame, 79, 22, 1, 1, 16, 1, 32, None, 2, 5, 6, 1, 5, 8, 8, 7]),
                  ('vgg13',         vgg13_bn,       [DataFrame, 50, 22, 1, 1, 1, 1, 2, None, 2, 5, 0, 0, 5, 1, 0, 5]),
                  ('alexnet',       alexnet,        [DataFrame, 28, 22, 1, 1, 1, 1, 2, None, 2, 3, 0, 0, 4, 1, 0, 5]),
                  ('squeezenet1_0', squeezenet1_0,  [DataFrame, 76, 22, 1, 1, 1, 1, 8, None, 2, 4, 0, 0, 5, 1, 8, 7]),
                  ('densenet161',   densenet161,    [DataFrame, 585, 22, 1, 1, 42, 1, 9, None, 2, 6, 7, 1, 5, 12, 78, 8]), # 11, 12 revise values after imorting cahnges from core.ipynb
                  ('xresnet34',     xresnet34,      [DataFrame, 219, 22, 1, 1, 71, 1, 120, None, 2, 5, 11, 1, 5, 8, 16, 8]) 
]

%%capture
for i, model in enumerate(models_to_test):
  _, m, res = model
  print(m)
  learn = cnn_learner(
      dls, 
      m, 
      opt_func=partial(Adam, lr=slice(3e-3), wd=0.01, eps=1e-8), 
      metrics=error_rate, 
      config=cnn_config(ps=0.33)).to_fp16()

  run_tests(ConvNav(learn, learn.summary()), i)

Tests