import os
import json
import math
import numpy as np
from PIL import Image

import torch
from torch.utils.data import Dataset, DataLoader, IterableDataset
import torchvision.transforms.functional as TF

import pytorch_lightning as pl

import datasets
from utils.misc import get_rank
import torch.nn.functional as F
from glob import glob

class BlenderDatasetBase():
    def setup(self, config, split):
        self.config = config
        self.split = split
        self.rank = get_rank()

        self.has_mask = True
        self.apply_mask = self.config.get('apply_mask', True)

        with open(os.path.join(self.config.root_dir, f"transforms_{self.split}.json"), 'r') as f:
            meta = json.load(f)
        scene_name = self.config.root_dir.split('/')[-1]
        if 'w' in meta and 'h' in meta:
            W, H = int(meta['w']), int(meta['h'])
        else:
            W, H = 800, 800

        if 'img_wh' in self.config:
            w, h = self.config.img_wh
            assert round(W / w * h) == H
        elif 'img_downscale' in self.config:
            w, h = W // self.config.img_downscale, H // self.config.img_downscale
        else:
            raise KeyError("Either img_wh or img_downscale should be specified.")
        
        self.w, self.h = w, h
        self.img_wh = (self.w, self.h)

        self.near, self.far = self.config.near_plane, self.config.far_plane

        self.focal = 0.5 * w / math.tan(0.5 * meta['camera_angle_x']) # scaled focal length
        self.fx, self.fy = self.focal, self.focal
        self.K = torch.tensor([[self.focal, 0, w/2, 0], [0, self.focal, h/2, 0], [0, 0, 1, 0]]).float().to(self.rank)

        # ray directions for all pixels, same for all images (same H, W, focal)

        if split == "test": 
            self.all_c2w, self.all_images, self.all_fg_masks = [], [], []
            for i, frame in enumerate(meta['frames']):
                c2w = torch.from_numpy(np.array(frame['transform_matrix'])[:3, :4])
                self.all_c2w.append(c2w)

                img_path = os.path.join(self.config.root_dir, f"{frame['file_path']}.png")
                if img_path.endswith(".png.png"):
                    img_path = img_path[:-4]
                img = Image.open(img_path)
                img = img.resize(self.img_wh, Image.BICUBIC)
                img = TF.to_tensor(img).permute(1, 2, 0) # (4, h, w) => (h, w, 4)

                self.all_fg_masks.append(img[..., -1]) # (h, w)
                self.all_images.append(img[...,:3])

            self.all_c2w, self.all_images, self.all_fg_masks = \
                torch.stack(self.all_c2w, dim=0).float().to(self.rank), \
                torch.stack(self.all_images, dim=0).float().to(self.rank), \
                torch.stack(self.all_fg_masks, dim=0).float().to(self.rank)

        elif split == "train" or split == "val": 
            self.all_c2w, self.all_images, self.all_fg_masks = [], [], []

            c2ws = f"./data/nbvpose/{scene_name}/*.npy"
            c2ws = glob(c2ws)
            c2ws.sort()

            for c2w in c2ws: 
                img_path = c2w.replace("nbvpose", "images").replace("npy", "png")
                #img_path = c2w.replace("c2w", "gt").replace("npy", "png")
                c2w = np.load(c2w)
                c2w = torch.from_numpy(c2w)[:3,:4]
                self.all_c2w.append(c2w)

                img = Image.open(img_path)
                img = img.resize(self.img_wh, Image.BICUBIC)
                img = TF.to_tensor(img).permute(1, 2, 0) # (4, h, w) => (h, w, 4)

                self.all_fg_masks.append(img[..., -1]) # (h, w)
                self.all_images.append(img[...,:3])

            self.all_c2w, self.all_images, self.all_fg_masks = \
                torch.stack(self.all_c2w, dim=0).float().to(self.rank), \
                torch.stack(self.all_images, dim=0).float().to(self.rank), \
                torch.stack(self.all_fg_masks, dim=0).float().to(self.rank)

    def add_one_image(self, nbv_c2w, img_path): 
        img = Image.open(img_path)
        img = img.resize(self.config.img_wh, Image.BICUBIC)
        img = TF.to_tensor(img).to(self.rank).permute(1, 2, 0).float() # (4, h, w) => (h, w, 4)

        image = img[..., :3] # (h, w)
        fg_mask = img[...,-1]

        nbv_c2w = nbv_c2w[:3, :4]

        self.all_c2w, self.all_images, self.all_fg_masks = \
            torch.cat([self.all_c2w, nbv_c2w[None, ...]], dim=0),  \
            torch.cat([self.all_images, image[None, ...]], dim=0),  \
            torch.cat([self.all_fg_masks, fg_mask[None, ...]], dim=0)

class BlenderDataset(Dataset, BlenderDatasetBase):
    def __init__(self, config, split):
        self.setup(config, split)

    def __len__(self):
        return len(self.all_images)
    
    def __getitem__(self, index):
        return {
            'index': index
        }


class BlenderIterableDataset(IterableDataset, BlenderDatasetBase):
    def __init__(self, config, split):
        self.setup(config, split)

    def __iter__(self):
        while True:
            yield {}


@datasets.register('blender_retrain')
class BlenderDataModule(pl.LightningDataModule):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.n_train_images = 100

    def setup(self, stage=None):
        if stage in [None, 'fit']:
            self.train_dataset = BlenderIterableDataset(self.config, self.config.train_split)
        if stage in [None, 'fit', 'validate']:
            self.val_dataset = BlenderDataset(self.config, self.config.val_split)
        if stage in [None, 'test']:
            self.test_dataset = BlenderDataset(self.config, self.config.test_split)
        if stage in [None, 'predict']:
            self.predict_dataset = BlenderDataset(self.config, self.config.train_split)

    def prepare_data(self):
        pass
    
    def general_loader(self, dataset, batch_size):
        sampler = None
        return DataLoader(
            dataset, 
            num_workers=32, 
            batch_size=batch_size,
            pin_memory=True,
            sampler=sampler
        )
    
    def train_dataloader(self):
        return self.general_loader(self.train_dataset, batch_size=1)

    def val_dataloader(self):
        return self.general_loader(self.val_dataset, batch_size=1)

    def test_dataloader(self):
        return self.general_loader(self.test_dataset, batch_size=1) 

    def predict_dataloader(self):
        return self.general_loader(self.predict_dataset, batch_size=1)       
