import sys
import os
import json
from PIL import Image
from torch import load, no_grad, sigmoid, device
from torch.nn import Sequential
from torchvision import transforms
from effdet import create_model


class DocumentClassifier():
    """This class uses effdet model artifact to process images
    and classify the image if there is evidence of photo manipulation."""
    def __init__(self, _device='cuda'):
        """ Loads model from directory"""
        self.device = device(_device)
        self.load_model()

    def load_model(self):
        """ Loading the model includes the effdet models used for feature generation."""
        # Load classifier model
        path = os.path.dirname(__file__)
        with open(path+"/config.json") as config_file:
            self.cfg = json.load(config_file)
        self.model = load(path+"/documentfraud.pytorch.model")
        # Load feature generator
        d1 = list(create_model('tf_efficientdet_d1', 
                                num_classes=self.cfg["N_EFFDET_CLASSES"],
                                pretrained=True).backbone.children())
        layers_d1 = d1[:2]
        for layer in d1[2][:3]:
            layers_d1.append(layer)
        self.effdet_generator = Sequential(*layers_d1).eval().to(self.device)
        self.edm = load('documentfraud_model/tf_efficientdet_d1.pytorch.model', map_location=self.device)
        self.edm.eval().to(self.device)
        self.edm.fpn.eval().to(self.device)
        transform_l = [transforms.Resize([640, 640]),
                transforms.ToTensor()]
        self.transform = transforms.Compose(transform_l)
        return
    
    def effdet_feature_extractor(self, img):
        '''Get the feature from efficient det'''
        with no_grad():
            fpn_out = self.effdet_generator(self.transform(img).unsqueeze(0).to(self.device))[0].unsqueeze(0)
        with no_grad():
            fpn_out = self.edm.backbone(fpn_out)
            fpn_out = self.edm.fpn(fpn_out)
        return fpn_out
    
    def get_classes(self, result):
        predicted_score = sigmoid(result).tolist()
        predicted_class = [0] * len(predicted_score)
        for flag in self.cfg["label_idxs"].keys():
            scr = predicted_score[self.cfg["label_idxs"][flag]]
            threshold = self.cfg["THRESHOLDS"][flag]
            predicted_class[self.cfg["label_idxs"][flag]] = 1 if scr >= threshold  else 0
        return predicted_class, predicted_score
    
    def classify_document(self, filename, image):
        """ Classify document image 
        Inputs:
            filename -- filename of the input image
            image -- a PIL image 

        Returns:

        """

        response = {'filename': filename}
        # Get feature from features folder if it exists
        # otherwise load from image
        fpn = self.effdet_feature_extractor(image)
        y_pred = self.model(fpn[self.cfg["FEATURE_NODE"]])
        predicted_classes, predicted_scores = self.get_classes(y_pred[0])
        response['classes'] = {}
        response['scores'] = {}
        for f in self.cfg["label_idxs"].keys():
            response['classes'][f] = predicted_classes[self.cfg["label_idxs"][f]]
            response['scores'][f] = predicted_scores[self.cfg["label_idxs"][f]]
        return response


if __name__ == "__main__":
    classifier = DocumentClassifier()
    image_filename = sys.argv[1]
    image = Image.open(image_filename).convert("RGB")
    result = classifier.classify_document(image_filename)
    print(result)

