import os
import torchvision.transforms as T
import cv2
from tqdm import tqdm
import detection.transforms as transforms
from dataloaders import get_direction

def generate_predictions_bilateral(model,device,testpath_,cor_dict,dset='aiims',preds_folder='preds_new'):
    transform = T.Compose([T.ToPILImage(),T.ToTensor()])
    model.eval()
    for label in ['mal','ben']:
        testpath = os.path.join(testpath_,label)
        # testpath = os.path.join(dataset_path,'Training', 'train',label)
        testimg = os.path.join(testpath, 'images')

        #preds_folder  = 'preds_new'
        os.makedirs(os.path.join(testpath, preds_folder),exist_ok=True)

        if not os.path.exists(os.path.join(testpath,preds_folder)):
            os.makedirs(os.path.join(testpath+preds_folder),exist_ok = True)

        for file in tqdm(os.listdir(testimg)):
            img1 = cv2.imread(os.path.join(testimg,file))
            img1 = transform(img1)
            # if False:
            if(os.path.join(testimg,file) in cor_dict and os.path.isfile(cor_dict[os.path.join(testimg,file)])):
                print('Using Bilateral')
                img2 = cv2.imread(cor_dict[os.path.join(testimg,file)])
                img2 = transform(img2)
                if(get_direction(dset,file)==1):
                    img1,_ = transforms.RandomHorizontalFlip(1.0)(img1)
                    
                    images = [img1.to(device),img2.to(device)]
                    output = model([images])[0]
                    img1,output = transforms.RandomHorizontalFlip(1.0)(img1,output)
                else:
                    img2,_ = transforms.RandomHorizontalFlip(1.0)(img2)
                    
                    images = [img1.to(device),img2.to(device)]
                    output = model([images])[0]
            else:
                print('Using FRCNN')
                output = model.frcnn([img1.to(device)])[0]
            #output = model.frcnn([img1.to(device)])[0]
            boxes = output['boxes']
            scores = output['scores']
            labels = output['labels']
            f = open(os.path.join(testpath,preds_folder,file[:-4]+'.txt'),'w')
            for i in range(len(boxes)):
                box = boxes[i].detach().cpu().numpy()
                #f.write('{} {} {} {} {} {}\n'.format(scores[i].item(),labels[i].item(),box[0],box[1],box[2],box[3]))
                f.write('{} {} {} {} {}\n'.format(scores[i].item(),box[0],box[1],box[2],box[3]))


def generate_predictions(model,device,testpath_,preds_folder='preds_frcnn'):
    transform = T.Compose([T.ToPILImage(),T.ToTensor()])
    model.eval()
    for label in ['mal','ben']:
        testpath = os.path.join(testpath_,label)
        # testpath = os.path.join(dataset_path,'Training', 'train',label)
        testimg = os.path.join(testpath, 'images')

        #preds_folder  = 'preds_new'
        os.makedirs(os.path.join(testpath, preds_folder),exist_ok=True)

        if not os.path.exists(os.path.join(testpath,preds_folder)):
            os.makedirs(os.path.join(testpath+preds_folder),exist_ok = True)

        for file in tqdm(os.listdir(testimg)):
            im = cv2.imread(os.path.join(testimg,file))
            if file == 'Mass-Training_P_00444_LEFT_CC.png':
                print('Test this')
                continue
            im = transform(im)

            output = model([im.to(device)])[0]
            boxes = output['boxes'] #/ FAC
            scores = output['scores']
            labels = output['labels']
            f = open(os.path.join(testpath,preds_folder,file[:-4]+'.txt'),'w')
            for i in range(len(boxes)):
                box = boxes[i].detach().cpu().numpy()
                f.write('{} {} {} {} {}\n'.format(scores[i].item(),box[0],box[1],box[2],box[3]))