# based on https://github.com/universome/fvd-comparison/blob/master/compare_models.py; huge thanks!
import os
import numpy as np
import io
import re
import requests
import html
import hashlib
import urllib
import urllib.request
import scipy.linalg
import multiprocessing as mp
import glob


from tqdm import tqdm
from typing import Any, List, Tuple, Union, Dict, Callable

from torchvision.io import read_video
import torch; torch.set_grad_enabled(False)
from einops import rearrange

from nitro.util import isvideo

def compute_frechet_distance(mu_sample,sigma_sample,mu_ref,sigma_ref) -> float:
    print('Calculate frechet distance...')
    m = np.square(mu_sample - mu_ref).sum()
    s, _ = scipy.linalg.sqrtm(np.dot(sigma_sample, sigma_ref), disp=False) # pylint: disable=no-member
    fid = np.real(m + np.trace(sigma_sample + sigma_ref - s * 2))

    return float(fid)


def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    mu = feats.mean(axis=0) # [d]
    sigma = np.cov(feats, rowvar=False) # [d, d]

    return mu, sigma


def open_url(url: str, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False) -> Any:
    """Download the given URL and return a binary-mode file object to access the data."""
    assert num_attempts >= 1

    # Doesn't look like an URL scheme so interpret it as a local filename.
    if not re.match('^[a-z]+://', url):
        return url if return_filename else open(url, "rb")

    # Handle file URLs.  This code handles unusual file:// patterns that
    # arise on Windows:
    #
    # file:///c:/foo.txt
    #
    # which would translate to a local '/c:/foo.txt' filename that's
    # invalid.  Drop the forward slash for such pathnames.
    #
    # If you touch this code path, you should test it on both Linux and
    # Windows.
    #
    # Some internet resources suggest using urllib.request.url2pathname() but
    # but that converts forward slashes to backslashes and this causes
    # its own set of problems.
    if url.startswith('file://'):
        filename = urllib.parse.urlparse(url).path
        if re.match(r'^/[a-zA-Z]:', filename):
            filename = filename[1:]
        return filename if return_filename else open(filename, "rb")

    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()

    # Download.
    url_name = None
    url_data = None
    with requests.Session() as session:
        if verbose:
            print("Downloading %s ..." % url, end="", flush=True)
        for attempts_left in reversed(range(num_attempts)):
            try:
                with session.get(url) as res:
                    res.raise_for_status()
                    if len(res.content) == 0:
                        raise IOError("No data received")

                    if len(res.content) < 8192:
                        content_str = res.content.decode("utf-8")
                        if "download_warning" in res.headers.get("Set-Cookie", ""):
                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
                            if len(links) == 1:
                                url = requests.compat.urljoin(url, links[0])
                                raise IOError("Google Drive virus checker nag")
                        if "Google Drive - Quota exceeded" in content_str:
                            raise IOError("Google Drive download quota exceeded -- please try again later")

                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
                    url_name = match[1] if match else url
                    url_data = res.content
                    if verbose:
                        print(" done")
                    break
            except KeyboardInterrupt:
                raise
            except:
                if not attempts_left:
                    if verbose:
                        print(" failed")
                    raise
                if verbose:
                    print(".", end="", flush=True)

    # Return data as file object.
    assert not return_filename
    return io.BytesIO(url_data)

def load_video(ip):
    vid, *_ = read_video(ip)
    vid = rearrange(vid, 't h w c -> t c h w').to(torch.uint8)
    return vid

def get_data_from_str(input_str,nprc = None):
    assert os.path.isdir(input_str), f'Specified input folder "{input_str}" is not a directory'
    vid_filelist = glob.glob(os.path.join(input_str,'*.mp4'))
    print(f'Found {len(vid_filelist)} videos in dir {input_str}')

    if nprc is None:
        try:
            nprc = mp.cpu_count()
        except NotImplementedError:
            print('WARNING: cpu_count() not avlailable, using only 1 cpu for video loading')
            nprc = 1

    pool = mp.Pool(processes=nprc)

    vids = []
    for v in tqdm(pool.imap_unordered(load_video,vid_filelist),total=len(vid_filelist),desc='Loading videos...'):
        vids.append(v)


    vids = torch.stack(vids,dim=0).float()

    return vids

def get_stats(stats):
    assert os.path.isfile(stats) and stats.endswith('.npz'), f'no stats found under {stats}'

    print(f'Using precomputed statistics under {stats}')
    stats = np.load(stats)
    stats = {key: stats[key] for key in stats.files}

    return stats




@torch.no_grad()
def compute_fvd(ref_input, sample_input, bs=32,
                ref_stats=None,
                sample_stats=None,
                nprc_load=None):



    calc_stats = ref_stats is None or sample_stats is None

    if calc_stats:

        only_ref = sample_stats is not None
        only_sample = ref_stats is not None


        if isinstance(ref_input,str) and not only_sample:
            ref_input = get_data_from_str(ref_input,nprc_load)

        if isinstance(sample_input, str) and not only_ref:
            sample_input = get_data_from_str(sample_input, nprc_load)

        stats = compute_statistics(sample_input,ref_input,
                                        device='cuda' if torch.cuda.is_available() else 'cpu',
                                        bs=bs,
                                        only_ref=only_ref,
                                        only_sample=only_sample)

        if only_ref:
            stats.update(get_stats(sample_stats))
        elif only_sample:
            stats.update(get_stats(ref_stats))



    else:
        stats = get_stats(sample_stats)
        stats.update(get_stats(ref_stats))

    fvd = compute_frechet_distance(**stats)

    return {'FVD' : fvd,}


@torch.no_grad()
def compute_statistics(videos_fake, videos_real, device: str='cuda', bs=32, only_ref=False,only_sample=False) -> Dict:
    detector_url = 'https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt?dl=1'
    detector_kwargs = dict(rescale=True, resize=True, return_features=True) # Return raw features before the softmax layer.

    with open_url(detector_url, verbose=False) as f:
        detector = torch.jit.load(f).eval().to(device)



    assert not (only_sample and only_ref), 'only_ref and only_sample arguments are mutually exclusive'

    ref_embed, sample_embed = [], []

    info = f'Computing I3D activations for FVD score with batch size {bs}'

    if only_ref:

        if not isvideo(videos_real):
            # if not is video we assume to have numpy arrays pf shape (n_vids, t, h, w, c) in range [0,255]
            videos_real = torch.from_numpy(videos_real).permute(0, 4, 1, 2, 3).float()
            print(videos_real.shape)

        if videos_real.shape[0] % bs == 0:
            n_secs = videos_real.shape[0] // bs
        else:
            n_secs = videos_real.shape[0] // bs + 1

        videos_real = torch.tensor_split(videos_real, n_secs, dim=0)

        for ref_v in tqdm(videos_real, total=len(videos_real),desc=info):

            feats_ref = detector(ref_v.to(device).contiguous(), **detector_kwargs).cpu().numpy()
            ref_embed.append(feats_ref)

    elif only_sample:

        if not isvideo(videos_fake):
            # if not is video we assume to have numpy arrays pf shape (n_vids, t, h, w, c) in range [0,255]
            videos_fake = torch.from_numpy(videos_fake).permute(0, 4, 1, 2, 3).float()
            print(videos_fake.shape)

        if videos_fake.shape[0] % bs == 0:
            n_secs = videos_fake.shape[0] // bs
        else:
            n_secs = videos_fake.shape[0] // bs + 1

        videos_real = torch.tensor_split(videos_real, n_secs, dim=0)

        for sample_v in tqdm(videos_fake, total=len(videos_real),desc=info):
            feats_sample = detector(sample_v.to(device).contiguous(), **detector_kwargs).cpu().numpy()
            sample_embed.append(feats_sample)


    else:

        if not isvideo(videos_real):
            # if not is video we assume to have numpy arrays pf shape (n_vids, t, h, w, c) in range [0,255]
            videos_real = torch.from_numpy(videos_real).permute(0, 4, 1, 2, 3).float()

        if not isvideo(videos_fake):
            videos_fake = torch.from_numpy(videos_fake).permute(0, 4, 1, 2, 3).float()

        if videos_fake.shape[0] % bs == 0:
            n_secs = videos_fake.shape[0] // bs
        else:
            n_secs = videos_fake.shape[0] // bs + 1

        videos_real = torch.tensor_split(videos_real, n_secs, dim=0)
        videos_fake = torch.tensor_split(videos_fake, n_secs, dim=0)

        for ref_v, sample_v in tqdm(zip(videos_real,videos_fake),total=len(videos_fake),desc=info):
            # print(ref_v.shape)
            # ref_v = torch.nn.functional.interpolate(ref_v, size=(sample_v.shape[2], 256, 256), mode='trilinear', align_corners=False)
            # sample_v = torch.nn.functional.interpolate(sample_v, size=(sample_v.shape[2], 256, 256), mode='trilinear', align_corners=False)


            feats_sample = detector(sample_v.to(device).contiguous(), **detector_kwargs).cpu().numpy()
            feats_ref = detector(ref_v.to(device).contiguous(), **detector_kwargs).cpu().numpy()
            sample_embed.append(feats_sample)
            ref_embed.append(feats_ref)

    out = dict()
    if len(sample_embed) > 0:
        sample_embed = np.concatenate(sample_embed,axis=0)
        mu_sample, sigma_sample = compute_stats(sample_embed)
        out.update({'mu_sample': mu_sample,
                    'sigma_sample': sigma_sample})

    if len(ref_embed) > 0:
        ref_embed = np.concatenate(ref_embed,axis=0)
        mu_ref, sigma_ref = compute_stats(ref_embed)
        out.update({'mu_ref': mu_ref,
                    'sigma_ref': sigma_ref})


    return out