import os
from os.path import join, exists
import h5py
import numpy as np

import torch
import torchtext
from torch.functional import F
from torchvision.transforms import ToTensor
from PIL import Image
import tarfile
import io
import os
import time

def iou(candidates, gt):
    start, end = candidates[:,0], candidates[:,1]
    s, e = gt[0].float(), gt[1].float()
    # print(s.dtype, start.dtype)
    inter = end.min(e) - start.max(s)
    union = end.max(e) - start.min(s)
    return inter.clamp(min=0) / union

def score2d_to_moments_scores(score2d, num_clips, duration):
    #grids = score2d.nonzero() # Will raise warning
    grids = torch.nonzero(score2d, as_tuple=False) 
    scores = score2d[grids[:,0], grids[:,1]]
    grids[:, 1] += 1
    moments = grids * duration / num_clips
    return moments, scores

def moment_to_iou2d(moment, num_clips, duration):
    iou2d = torch.ones(num_clips, num_clips)
    candidates, _ = score2d_to_moments_scores(iou2d, num_clips, duration)
    iou2d = iou(candidates, moment).reshape(num_clips, num_clips)
    return iou2d

def avgfeats(feats, num_pre_clips):
    # Produce the feature of per video into fixed shape (e.g. 256*4096)
    # Input Example: feats (torch.tensor, ?x4096); num_pre_clips (256)
    num_src_clips = feats.size(0)
    idxs = torch.arange(0, num_pre_clips+1, 1.0) / num_pre_clips * num_src_clips
    idxs = idxs.round().long().clamp(max=num_src_clips-1)
    # To prevent a empty selection, check the idxs
    meanfeats = []
    for i in range(num_pre_clips):
        s, e = idxs[i], idxs[i+1]
        if s < e:
            meanfeats.append(feats[s:e].mean(dim=0))
        else:
            meanfeats.append(feats[s])
    return torch.stack(meanfeats)
    
def video2feats(feat_file, vids, num_pre_clips, dataset_name):
    assert exists(feat_file)
    vid_feats = {}
    with h5py.File(feat_file, 'r') as f:
        for vid in vids:
            if dataset_name == "activitynet":
                feat = f[vid]['c3d_features'][:]
            else:
                feat = f[vid][:]
            feat = F.normalize(torch.from_numpy(feat),dim=1)
            vid_feats[vid] = avgfeats(feat, num_pre_clips) 
    return vid_feats

def video2clips(vid_root, vids, trans, nb_frames=16, sample_interval=32):
    #vids = ['v_QOlSCBRmfWY', 'v_ehGHCYKzyZ8']
    vid_clips = {}
    for vid in vids: # list video name
        stime = time.time()
        stackframes = []
        vid_path = os.path.join(vid_root, "{}.tar".format(vid))
        tar = tarfile.open(vid_path)
        #for member in tar.getmembers(): # return tarInfo
        for member in tar.getnames(): # return data path
            frmIdx = int(member.split('/')[-1].split('.')[0])
            if frmIdx % sample_interval != 0:
                continue
            f = tar.extractfile(member)
            image = f.read()
            image = Image.open(io.BytesIO(image))
            image = trans(image)
            stackframes.append(image)
        #vid_clips[vid] = torch.stack(stackframes).permute(1, 0, 2, 3)
        # Padding to valid length
        if len(stackframes) < nb_frames:
            stackframes.extend([stackframes[0]] * (nb_frames - len(stackframes)))
        total_frames = len(stackframes)
        valid_frames = total_frames // nb_frames * nb_frames
        stackframes = torch.stack(stackframes[:valid_frames])
        vid_clips[vid] = torch.stack(stackframes.chunk(valid_frames//nb_frames, dim=0)) # size: [num_clip, nb, 3, H, W]
        vid_clips[vid] = vid_clips[vid].permute(0, 2, 1, 3, 4)
        etime = time.time()
        print("process video {} using {}s".format(vid, etime-stime))
    return vid_clips

def frm2clips(frames, nb_frames=16):
    """frames: frame list; nb_frames: length of each clip"""
    total_frames = len(frames)
    valid_frames = total_frames // nb_frames * nb_frames
    videoClip = []
    for i in range(valid_frames // nb_frames):
        clip = []
        for j in range(i*nb_frames+1, (i+1)*nb_frames+1):
            clip.append(frames[j])
        clip = torch.stack(clip) # size: [nb, 3, H, W]
        videoClip.append(clip)
    return torch.stack(videoClip) # size: [num_clip, nb, 3, H, W]

def embedding(sentence, vocabs=[], embedders=[]):
    if len(vocabs) == 0:
        vocab = torchtext.vocab.pretrained_aliases["glove.6B.300d"]()
        vocab.itos.extend(['<unk>'])
        vocab.stoi['<unk>'] = vocab.vectors.shape[0]
        vocab.vectors = torch.cat(
            [vocab.vectors, torch.zeros(1, vocab.dim)],
            dim=0
        )
        vocabs.append(vocab)
    
    if len(embedders) == 0:
        embedder = torch.nn.Embedding.from_pretrained(vocab.vectors)
        embedders.append(embedder)
    
    vocab, embedder = vocabs[0], embedders[0]
    word_idxs = torch.tensor([vocab.stoi.get(w.lower(), 400000) \
        for w in sentence.split()], dtype=torch.long)
    return embedder(word_idxs)
