import io
import numpy as np
import test_utils
import unittest

from cvmodelz.classifiers import Classifier
from cvmodelz.models import ModelFactory

class ClassifierLoadingTests(unittest.TestCase):

	def new_clf(self, key, **kwargs):
		model = ModelFactory.new(key, pretrained_model=None)
		return model, Classifier(model, **kwargs)

	def load_model(self, key, finetune):
		init_cls = 1000 if finetune else 200
		final_cls = 200

		model = ModelFactory.new(key, pretrained_model=None, n_classes=init_cls)
		_, clf = self.new_clf(key)

		self.assertTrue(*test_utils.is_any_different(model, clf.model))

		with test_utils.memory_file() as f:
			model.save(f)
			clf.load(f, n_classes=final_cls, finetune=finetune)

		""" if finetune is True, then the shapes of the classification
		layer differ, hence, strict should be False """
		self.assertTrue(*test_utils.is_all_equal(model, clf.model, strict=not finetune))

	def load_for_finetune(self, key):
		self.load_model(key, finetune=True)

	def load_for_inference(self, key):
		self.load_model(key, finetune=False)

test_utils.add_tests(ClassifierLoadingTests.load_for_finetune,
	model_list=ModelFactory.get_models(["cvmodelz", "chainercv2"]))

test_utils.add_tests(ClassifierLoadingTests.load_for_inference,
	model_list=ModelFactory.get_models(["cvmodelz", "chainercv2"]))