
from numpy.lib.npyio import save
from torchvision import models
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torch.hub
import torch.fft
import torchvision.transforms as T
from torchvision.utils import save_image
import torch.nn.functional as F
import numpy as np
import math
import pandas as pd
from torchvision.transforms.functional import normalize
from dataset import ImageFolderWithPaths
import os

dlab = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()

def get_train_loader(batch_size=1, test_mean=False):
    """get the train loader"""
    data_folder = "/data/movie-associations/train" 

    # Use ImageNet mean and std for segmentation
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    trf = T.Compose([
                    T.Resize(256),
                    T.CenterCrop(224),
                    T.ToTensor()
                    ])

    dataset = ImageFolderWithPaths(data_folder, transform=trf)
    dataloader = torch.utils.data.DataLoader(dataset,
                                            batch_size=batch_size,
                                            num_workers=0,
                                            pin_memory=True,
                                            shuffle = False)
    if not test_mean:
        return dataloader
    
    mean = 0.0
    for images, _, _ in dataloader:
        batch_samples = images.size(0) 
        images = images.view(batch_samples, images.size(1), -1)
        mean += images.mean(2).sum(0)
    mean = mean / len(dataloader.dataset)

    var = 0.0
    for images, _, _ in dataloader:
        batch_samples = images.size(0)
        images = images.view(batch_samples, images.size(1), -1)
        var += ((images - mean.unsqueeze(1))**2).sum([0,2])
    std = torch.sqrt(var / (len(dataloader.dataset)*224*224))
    
    print(f'mean = {mean} \t std = {std}')

    return dataloader

class NormalizeInverse(T.Normalize):
  """
  Undoes the normalization and returns the reconstructed images in the input domain.
  """

  def __init__(self, mean, std):
      mean = torch.as_tensor(mean)
      std = torch.as_tensor(std)
      std_inv = 1 / (std + 1e-7)
      mean_inv = -mean * std_inv
      super().__init__(mean=mean_inv, std=std_inv)

  def __call__(self, tensor):
      return super().__call__(tensor.clone())    

def segment_vals(inp, device, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], net=dlab, remove='background'):
    """ 
    Function to take a tensor (from dataloader in get_activations) and return as segmented tensor
    Segmentation uses Google DeepLab v3 which requires normalization using imagenet mean and std
    
    inp: is torch.Size([1,3,h,w]) or whatever h*w are
    """
    #normalize input to imgnet for dlab
    normalize = T.Normalize(mean=mean , std=std)
    inp_norm = normalize(inp)

    #pass through and get segmentation maps
    net.to(device)
    out = net(inp_norm)['out']
    segmap = torch.argmax(out.squeeze(), dim=0)
    segmap = segmap.float()

    #place objects on solid background
    dataset_mean = torch.Tensor([0.2629, 0.2301, 0.2021]).to(device)
    r = torch.where(segmap != 0, inp[0,0,:,:], dataset_mean[0]) #change this to movie dataset mean
    g = torch.where(segmap != 0, inp[0,1,:,:], dataset_mean[1])
    b = torch.where(segmap != 0, inp[0,2,:,:], dataset_mean[2])

    obj_img = torch.stack([r, g, b], axis=0)

    return obj_img

def main():
    loader = get_train_loader()
    _working_mov = ''
    for idx, (inputs, _, img_path) in enumerate(loader):

        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        if device.type == 'cpu':
            print('Using device:', device)

        img_path = img_path[0].split('/')
        mov = img_path[-2]
        if not mov == _working_mov:
            print(f'... working on {mov}')
            _working_mov = mov
        filename = img_path[-1]

        mov_dir = f'/data/movie-associations/train_segmented/{mov}'
        if not os.path.isdir(mov_dir):
            os.makedirs(mov_dir)
        save_path = f'{mov_dir}/segmented_{filename}'

        #if os.path.exists(save_path):
        #    continue

        inputs = inputs.float().to(device)
        segmented_img = segment_vals(inputs, device)

        save_image(segmented_img , save_path)

if __name__ == "__main__": 
    main()