from PIL import Image
import torch
import torchvision.transforms as T
from torchvision import models
import numpy as np
import OpenEXR
import Imath


fcn = None

def getRotoModel():
    global fcn
    fcn = models.segmentation.fcn_resnet101(pretrained=True).eval()

# Define the helper function
def decode_segmap(image, nc=21):

    label_colors = np.array([(0, 0, 0),  # 0=background
                           # 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle
               (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128),
               # 6=bus, 7=car, 8=cat, 9=chair, 10=cow
               (0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0),
               # 11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person
               (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128),
               # 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor
               (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128)])

    r = np.zeros_like(image).astype(np.uint8)
    g = np.zeros_like(image).astype(np.uint8)
    b = np.zeros_like(image).astype(np.uint8)

    for l in range(0, nc):
        idx = image == l
        r[idx] = label_colors[l, 0]
        g[idx] = label_colors[l, 1]
        b[idx] = label_colors[l, 2]

    rgb = np.stack([r, g, b], axis=2)
    return rgb

def createMatte(filename, matteName, size):
    # If .exr file, open like this
    if filename.endswith(".exr"):
        file = OpenEXR.InputFile(filename)
        pt = Imath.PixelType(Imath.PixelType.FLOAT)
        dw = file.header()['dataWindow']
        exrSize = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1)

        rgbf = [Image.frombytes("F", exrSize, file.channel(c, pt)) for c in "RGB"]

        extrema = [im.getextrema() for im in rgbf]
        darkest = min([lo for (lo,hi) in extrema])
        lighest = max([hi for (lo,hi) in extrema])
        scale = 255 / (lighest - darkest)
        def normalize_0_255(v):
            return (v * scale) + darkest
        rgb8 = [im.point(normalize_0_255).convert("L") for im in rgbf]
        img = Image.merge("RGB", rgb8)

    else:
        img = Image.open(filename).convert('RGB')

    trf = T.Compose([T.Resize(size),
                     T.ToTensor(), 
                     T.Normalize(mean = [0.485, 0.456, 0.406], 
                                 std = [0.229, 0.224, 0.225])])
    inp = trf(img).unsqueeze(0)
    if (fcn == None): getRotoModel()
    out = fcn(inp)['out']
    om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()  
    rgb = decode_segmap(om)
    im = Image.fromarray(rgb)
    im.save(matteName)