import pytorch_lightning as pl
import numpy as np
from PIL import Image
from src.ss.datasets_signboard_detection.dataset import Labelizer

class Color_convert():
    def __init__(self):
        super().__init__()
        self.labels = {'bien': "red"}
        
    def transform(self, label):
        return self.labels[label]
    
    def num_classes(self):
        return len(self.labels)

def compose(output, mask):
    w,h = mask.shape

    for i in range(0, w):
        for j in range(0,h):
            if (mask[i,j] > 0.5):
                output[i,j] = 1
    return output

class POIDetectionTask(pl.LightningModule):
    def __init__(self,
                 model,
                 score):
        super().__init__()
        
        self.model = model
        self.output = []
        self.score = score
        self.labelizer = Labelizer()
        self.color_convert = Color_convert()

    def forward(self, x):
        output = self.model(x)
        return output

    def predict_step(self, test_batch, batch_idx):
        images, targets = test_batch
        outputs = self(images)
        for target in outputs:
            shape = target['boxes']
            masks = target['masks']
            scores = target['scores']
            labels = target['labels']
            shape = shape.cpu().numpy()
            masks = masks.cpu().numpy()
            scores = scores.cpu().numpy()
            labels = labels.cpu().numpy()
            select_shape = []
            select_masks = []
            select_scores = []
            select_labels = []
            for i in range(len(scores)):
                if (scores[i]>self.score):
                    select_shape.append(shape[i])
                    select_masks.append(masks[i])
                    select_scores.append(scores[i])
                    select_labels.append(labels[i])

            output = {
                'boxes': np.array(select_shape, dtype=np.int32).tolist(),
                # 'masks': np.array(select_masks, dtype=np.uint8).tolist(),
                'scores': np.array(select_scores, dtype=np.float32).tolist(),
                'labels': np.array(select_labels, dtype=np.int32).tolist()
            }
            
            self.output.append(output)
    
    def on_predict_end(self):        
        self.output = self.output